@@ -0,0 +1,35 @@
|
||||
using System.Net;
|
||||
using System.Threading.Channels;
|
||||
|
||||
namespace ReverseLlama.Client;
|
||||
|
||||
internal sealed class ChannelHttpContent : HttpContent
|
||||
{
|
||||
private readonly CancellationToken _cancellationToken;
|
||||
private readonly ChannelReader<byte[]> _reader;
|
||||
|
||||
public ChannelHttpContent(ChannelReader<byte[]> reader, CancellationToken cancellationToken)
|
||||
{
|
||||
_reader = reader;
|
||||
_cancellationToken = cancellationToken;
|
||||
}
|
||||
|
||||
protected override Task SerializeToStreamAsync(Stream stream, TransportContext? context) =>
|
||||
SerializeToStreamAsync(stream, context, _cancellationToken);
|
||||
|
||||
protected override async Task SerializeToStreamAsync(Stream stream, TransportContext? context, CancellationToken cancellationToken)
|
||||
{
|
||||
using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(_cancellationToken, cancellationToken);
|
||||
|
||||
await foreach (var chunk in _reader.ReadAllAsync(linkedCts.Token))
|
||||
{
|
||||
await stream.WriteAsync(chunk, linkedCts.Token);
|
||||
}
|
||||
}
|
||||
|
||||
protected override bool TryComputeLength(out long length)
|
||||
{
|
||||
length = -1;
|
||||
return false;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
using ReverseLlama.Protocol;
|
||||
|
||||
namespace ReverseLlama.Client;
|
||||
|
||||
internal sealed class ClientOptions
|
||||
{
|
||||
public Uri Server { get; init; } = new("http://localhost:5001");
|
||||
|
||||
public Uri Upstream { get; init; } = new("http://localhost:11434");
|
||||
|
||||
public string TunnelPath { get; init; } = ProtocolConstants.DefaultTunnelPath;
|
||||
|
||||
public string? Token { get; init; }
|
||||
|
||||
public string ClientId { get; init; } = Environment.MachineName.ToLowerInvariant();
|
||||
|
||||
public TimeSpan ReconnectDelay { get; init; } = TimeSpan.FromSeconds(5);
|
||||
|
||||
public int ChunkSize { get; init; } = 64 * 1024;
|
||||
|
||||
public Uri TunnelUri
|
||||
{
|
||||
get
|
||||
{
|
||||
var builder = new UriBuilder(Server);
|
||||
builder.Scheme = builder.Scheme.ToLowerInvariant() switch
|
||||
{
|
||||
"http" => "ws",
|
||||
"https" => "wss",
|
||||
"ws" => "ws",
|
||||
"wss" => "wss",
|
||||
var unsupported => throw new InvalidOperationException($"Unsupported server URI scheme '{unsupported}'. Use http, https, ws, or wss.")
|
||||
};
|
||||
|
||||
if (string.IsNullOrWhiteSpace(builder.Path) || builder.Path == "/")
|
||||
{
|
||||
builder.Path = NormalizePath(TunnelPath);
|
||||
}
|
||||
|
||||
return builder.Uri;
|
||||
}
|
||||
}
|
||||
|
||||
public static ClientOptions Parse(string[] args)
|
||||
{
|
||||
var values = ParseArgs(args);
|
||||
|
||||
return new ClientOptions
|
||||
{
|
||||
Server = ReadUri(values, "server", "REVERSE_LLAMA_SERVER", "http://localhost:5001"),
|
||||
Upstream = ReadUri(values, "upstream", "REVERSE_LLAMA_UPSTREAM", "http://localhost:11434"),
|
||||
TunnelPath = NormalizePath(Read(values, "tunnel-path", "REVERSE_LLAMA_TUNNEL_PATH") ?? ProtocolConstants.DefaultTunnelPath),
|
||||
Token = Read(values, "token", "REVERSE_LLAMA_TOKEN"),
|
||||
ClientId = Read(values, "client-id", "REVERSE_LLAMA_CLIENT_ID") ?? Environment.MachineName.ToLowerInvariant(),
|
||||
ReconnectDelay = TimeSpan.FromSeconds(ReadInt(values, 5, "reconnect-delay", "REVERSE_LLAMA_RECONNECT_DELAY_SECONDS")),
|
||||
ChunkSize = ReadInt(values, 64 * 1024, "chunk-size", "REVERSE_LLAMA_CHUNK_SIZE")
|
||||
};
|
||||
}
|
||||
|
||||
public static string Usage =>
|
||||
"""
|
||||
ReverseLlama.Client options:
|
||||
--server <url> Server base URL, e.g. http://my-server:5050
|
||||
--upstream <url> Local upstream URL, e.g. http://localhost:11434
|
||||
--token <value> Optional token matching the server
|
||||
--client-id <name> Identifies this machine on the server; defaults to the machine name
|
||||
--tunnel-path <path> Defaults to /_reverse-llama/tunnel
|
||||
--reconnect-delay <sec> Defaults to 5
|
||||
--chunk-size <bytes> Defaults to 65536
|
||||
""";
|
||||
|
||||
private static Dictionary<string, string> ParseArgs(string[] args)
|
||||
{
|
||||
var values = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase);
|
||||
|
||||
for (var i = 0; i < args.Length; i++)
|
||||
{
|
||||
var arg = args[i];
|
||||
if (!arg.StartsWith("--", StringComparison.Ordinal))
|
||||
{
|
||||
throw new ArgumentException($"Unexpected argument '{arg}'.");
|
||||
}
|
||||
|
||||
var keyValue = arg[2..].Split('=', 2);
|
||||
if (keyValue.Length == 2)
|
||||
{
|
||||
values[keyValue[0]] = keyValue[1];
|
||||
continue;
|
||||
}
|
||||
|
||||
if (i + 1 >= args.Length || args[i + 1].StartsWith("--", StringComparison.Ordinal))
|
||||
{
|
||||
throw new ArgumentException($"Missing value for '{arg}'.");
|
||||
}
|
||||
|
||||
values[keyValue[0]] = args[++i];
|
||||
}
|
||||
|
||||
return values;
|
||||
}
|
||||
|
||||
private static string? Read(Dictionary<string, string> values, params string[] keys)
|
||||
{
|
||||
foreach (var key in keys)
|
||||
{
|
||||
if (values.TryGetValue(key, out var value) && !string.IsNullOrWhiteSpace(value))
|
||||
{
|
||||
return value;
|
||||
}
|
||||
|
||||
value = Environment.GetEnvironmentVariable(key);
|
||||
if (!string.IsNullOrWhiteSpace(value))
|
||||
{
|
||||
return value;
|
||||
}
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
private static int ReadInt(Dictionary<string, string> values, int fallback, params string[] keys)
|
||||
{
|
||||
var value = Read(values, keys);
|
||||
return int.TryParse(value, out var parsed) && parsed > 0 ? parsed : fallback;
|
||||
}
|
||||
|
||||
private static Uri ReadUri(Dictionary<string, string> values, string key, string envKey, string fallback)
|
||||
{
|
||||
var value = Read(values, key, envKey) ?? fallback;
|
||||
|
||||
if (!Uri.TryCreate(value, UriKind.Absolute, out var uri))
|
||||
{
|
||||
throw new ArgumentException($"'{value}' is not an absolute URI.");
|
||||
}
|
||||
|
||||
return uri;
|
||||
}
|
||||
|
||||
private static string NormalizePath(string path) =>
|
||||
path.StartsWith('/') ? path : $"/{path}";
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
using Microsoft.Extensions.DependencyInjection;
|
||||
using Microsoft.Extensions.Hosting;
|
||||
using Microsoft.Extensions.Logging;
|
||||
using ReverseLlama.Client;
|
||||
|
||||
try
|
||||
{
|
||||
var options = ClientOptions.Parse(args);
|
||||
|
||||
Console.WriteLine("ReverseLlama client");
|
||||
Console.WriteLine($" client id: {options.ClientId}");
|
||||
Console.WriteLine($" server tunnel: {options.TunnelUri}");
|
||||
Console.WriteLine($" local upstream: {options.Upstream}");
|
||||
|
||||
// Args are parsed by ClientOptions; keep them away from the host configuration.
|
||||
var builder = Host.CreateApplicationBuilder(new HostApplicationBuilderSettings { Args = [] });
|
||||
builder.Services.AddSingleton(options);
|
||||
builder.Services.AddSingleton<TunnelClient>();
|
||||
builder.Services.AddHostedService<TunnelWorker>();
|
||||
builder.Services.AddWindowsService(service => service.ServiceName = "ReverseLlamaClient");
|
||||
// The EventLog provider defaults to Warning; connection state is worth seeing there.
|
||||
builder.Logging.AddFilter<Microsoft.Extensions.Logging.EventLog.EventLogLoggerProvider>("ReverseLlama.Client", LogLevel.Information);
|
||||
|
||||
await builder.Build().RunAsync();
|
||||
return 0;
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
Console.Error.WriteLine(exception.Message);
|
||||
Console.Error.WriteLine();
|
||||
Console.Error.WriteLine(ClientOptions.Usage);
|
||||
return 1;
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
using System.Runtime.CompilerServices;
|
||||
|
||||
[assembly: InternalsVisibleTo("ReverseLlama.Client.Tests")]
|
||||
@@ -0,0 +1,18 @@
|
||||
<Project Sdk="Microsoft.NET.Sdk">
|
||||
|
||||
<ItemGroup>
|
||||
<ProjectReference Include="..\ReverseLlama.Protocol\ReverseLlama.Protocol.csproj" />
|
||||
</ItemGroup>
|
||||
|
||||
<ItemGroup>
|
||||
<PackageReference Include="Microsoft.Extensions.Hosting.WindowsServices" Version="8.0.1" />
|
||||
</ItemGroup>
|
||||
|
||||
<PropertyGroup>
|
||||
<OutputType>Exe</OutputType>
|
||||
<TargetFramework>net10.0</TargetFramework>
|
||||
<ImplicitUsings>enable</ImplicitUsings>
|
||||
<Nullable>enable</Nullable>
|
||||
</PropertyGroup>
|
||||
|
||||
</Project>
|
||||
@@ -0,0 +1,540 @@
|
||||
using System.Collections.Concurrent;
|
||||
using System.Net.WebSockets;
|
||||
using System.Text;
|
||||
using System.Text.Json;
|
||||
using Microsoft.Extensions.Logging;
|
||||
using Microsoft.Extensions.Logging.Abstractions;
|
||||
using ReverseLlama.Protocol;
|
||||
|
||||
namespace ReverseLlama.Client;
|
||||
|
||||
internal sealed class TunnelClient
|
||||
{
|
||||
private static readonly TimeSpan ModelRefreshInterval = TimeSpan.FromSeconds(15);
|
||||
private static readonly TimeSpan ModelRefreshTimeout = TimeSpan.FromSeconds(10);
|
||||
private static readonly JsonSerializerOptions JsonOptions = new(JsonSerializerDefaults.Web);
|
||||
private const string EmbeddingWarmupInput = "ReverseLlama warmup";
|
||||
|
||||
private readonly ConcurrentDictionary<string, UpstreamRequest> _activeRequests = new();
|
||||
private readonly HttpClient _httpClient;
|
||||
private readonly ClientOptions _options;
|
||||
private readonly ILogger<TunnelClient> _logger;
|
||||
private readonly object _modelSnapshotLock = new();
|
||||
private readonly SemaphoreSlim _sendLock = new(1, 1);
|
||||
private List<string> _lastActiveModels = [];
|
||||
private List<string> _lastModels = [];
|
||||
|
||||
public TunnelClient(ClientOptions options, ILogger<TunnelClient>? logger = null)
|
||||
{
|
||||
_options = options;
|
||||
_logger = logger ?? NullLogger<TunnelClient>.Instance;
|
||||
_httpClient = new HttpClient
|
||||
{
|
||||
Timeout = Timeout.InfiniteTimeSpan
|
||||
};
|
||||
}
|
||||
|
||||
public async Task RunAsync(CancellationToken cancellationToken)
|
||||
{
|
||||
while (!cancellationToken.IsCancellationRequested)
|
||||
{
|
||||
using var socket = new ClientWebSocket();
|
||||
socket.Options.KeepAliveInterval = TimeSpan.FromSeconds(30);
|
||||
|
||||
if (!string.IsNullOrWhiteSpace(_options.Token))
|
||||
{
|
||||
socket.Options.SetRequestHeader(ProtocolConstants.TokenHeader, _options.Token);
|
||||
}
|
||||
|
||||
socket.Options.SetRequestHeader(ProtocolConstants.ClientIdHeader, _options.ClientId);
|
||||
|
||||
using var connectionCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
|
||||
Task? modelRefreshTask = null;
|
||||
|
||||
try
|
||||
{
|
||||
_logger.LogInformation("Connecting to {TunnelUri}...", _options.TunnelUri);
|
||||
await socket.ConnectAsync(_options.TunnelUri, cancellationToken);
|
||||
_logger.LogInformation("Tunnel connected.");
|
||||
|
||||
modelRefreshTask = RefreshModelsLoopAsync(socket, connectionCts.Token);
|
||||
await ReceiveLoopAsync(socket, connectionCts.Token);
|
||||
|
||||
if (socket.CloseStatusDescription == ProtocolConstants.ReplacedCloseDescription)
|
||||
{
|
||||
_logger.LogWarning("This client was replaced by a newer tunnel client. Exiting.");
|
||||
return;
|
||||
}
|
||||
|
||||
_logger.LogInformation("Tunnel closed by server ({Reason}).", socket.CloseStatusDescription ?? "no reason given");
|
||||
}
|
||||
catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
|
||||
{
|
||||
break;
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
_logger.LogWarning("Tunnel disconnected: {Message}", exception.Message);
|
||||
}
|
||||
finally
|
||||
{
|
||||
connectionCts.Cancel();
|
||||
if (modelRefreshTask is not null)
|
||||
{
|
||||
try
|
||||
{
|
||||
await modelRefreshTask;
|
||||
}
|
||||
catch (OperationCanceledException)
|
||||
{
|
||||
}
|
||||
}
|
||||
|
||||
CancelAllActiveRequests();
|
||||
}
|
||||
|
||||
if (!cancellationToken.IsCancellationRequested)
|
||||
{
|
||||
_logger.LogInformation("Reconnecting in {Seconds:0.#} seconds...", _options.ReconnectDelay.TotalSeconds);
|
||||
await Task.Delay(_options.ReconnectDelay, cancellationToken);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private async Task RefreshModelsLoopAsync(ClientWebSocket socket, CancellationToken cancellationToken)
|
||||
{
|
||||
while (socket.State == WebSocketState.Open && !cancellationToken.IsCancellationRequested)
|
||||
{
|
||||
try
|
||||
{
|
||||
await RefreshModelsOnceAsync(socket, cancellationToken);
|
||||
}
|
||||
catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
|
||||
{
|
||||
break;
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
_logger.LogWarning("Failed to report upstream model list: {Message}", exception.Message);
|
||||
}
|
||||
|
||||
await Task.Delay(ModelRefreshInterval, cancellationToken);
|
||||
}
|
||||
}
|
||||
|
||||
private async Task RefreshModelsOnceAsync(ClientWebSocket socket, CancellationToken cancellationToken)
|
||||
{
|
||||
var modelsTask = TryRefreshModelListAsync(GetUpstreamModelsAsync, "listed", cancellationToken);
|
||||
var activeModelsTask = TryRefreshModelListAsync(GetActiveUpstreamModelsAsync, "active", cancellationToken);
|
||||
|
||||
await Task.WhenAll(modelsTask, activeModelsTask);
|
||||
|
||||
var snapshot = UpdateCachedModelSnapshot(modelsTask.Result, activeModelsTask.Result);
|
||||
|
||||
await SendAsync(
|
||||
socket,
|
||||
new TunnelMessage
|
||||
{
|
||||
Type = TunnelMessageTypes.ModelSnapshot,
|
||||
Models = snapshot.Models,
|
||||
ActiveModels = snapshot.ActiveModels
|
||||
},
|
||||
cancellationToken);
|
||||
|
||||
_logger.LogInformation(
|
||||
"Reported {ModelCount} listed and {ActiveModelCount} active upstream model(s).",
|
||||
snapshot.Models.Count,
|
||||
snapshot.ActiveModels.Count);
|
||||
}
|
||||
|
||||
private async Task<List<string>> GetUpstreamModelsAsync(CancellationToken cancellationToken)
|
||||
{
|
||||
using var request = new HttpRequestMessage(HttpMethod.Get, new Uri(_options.Upstream, "/api/tags"));
|
||||
using var response = await _httpClient.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, cancellationToken);
|
||||
response.EnsureSuccessStatusCode();
|
||||
|
||||
await using var stream = await response.Content.ReadAsStreamAsync(cancellationToken);
|
||||
using var document = await JsonDocument.ParseAsync(stream, cancellationToken: cancellationToken);
|
||||
|
||||
return ExtractModelNames(document.RootElement);
|
||||
}
|
||||
|
||||
private async Task<List<string>> GetActiveUpstreamModelsAsync(CancellationToken cancellationToken)
|
||||
{
|
||||
using var request = new HttpRequestMessage(HttpMethod.Get, new Uri(_options.Upstream, "/api/ps"));
|
||||
using var response = await _httpClient.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, cancellationToken);
|
||||
|
||||
if (!response.IsSuccessStatusCode)
|
||||
{
|
||||
_logger.LogDebug("Ollama /api/ps returned {StatusCode}; active model list will be empty.", response.StatusCode);
|
||||
return [];
|
||||
}
|
||||
|
||||
await using var stream = await response.Content.ReadAsStreamAsync(cancellationToken);
|
||||
using var document = await JsonDocument.ParseAsync(stream, cancellationToken: cancellationToken);
|
||||
|
||||
return ExtractModelNames(document.RootElement);
|
||||
}
|
||||
|
||||
internal static List<string> ExtractModelNames(JsonElement root)
|
||||
{
|
||||
var models = new List<string>();
|
||||
|
||||
if (root.ValueKind == JsonValueKind.Object
|
||||
&& root.TryGetProperty("models", out var ollamaModels)
|
||||
&& ollamaModels.ValueKind == JsonValueKind.Array)
|
||||
{
|
||||
AddModelNames(models, ollamaModels, "name");
|
||||
AddModelNames(models, ollamaModels, "model");
|
||||
}
|
||||
|
||||
return NormalizeModelNames(models);
|
||||
}
|
||||
|
||||
private static void AddModelNames(List<string> models, JsonElement array, string propertyName)
|
||||
{
|
||||
foreach (var item in array.EnumerateArray())
|
||||
{
|
||||
if (item.ValueKind == JsonValueKind.Object
|
||||
&& item.TryGetProperty(propertyName, out var model)
|
||||
&& model.ValueKind == JsonValueKind.String
|
||||
&& !string.IsNullOrWhiteSpace(model.GetString()))
|
||||
{
|
||||
models.Add(model.GetString()!);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private async Task<List<string>?> TryRefreshModelListAsync(
|
||||
Func<CancellationToken, Task<List<string>>> refresh,
|
||||
string listName,
|
||||
CancellationToken cancellationToken)
|
||||
{
|
||||
try
|
||||
{
|
||||
using var refreshCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
|
||||
refreshCts.CancelAfter(ModelRefreshTimeout);
|
||||
|
||||
return await refresh(refreshCts.Token);
|
||||
}
|
||||
catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
|
||||
{
|
||||
throw;
|
||||
}
|
||||
catch (OperationCanceledException)
|
||||
{
|
||||
_logger.LogWarning("Timed out refreshing {ModelListName} upstream model list.", listName);
|
||||
return null;
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
_logger.LogWarning(
|
||||
"Failed to refresh {ModelListName} upstream model list: {Message}",
|
||||
listName,
|
||||
exception.Message);
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
private (List<string> Models, List<string> ActiveModels) UpdateCachedModelSnapshot(
|
||||
List<string>? models,
|
||||
List<string>? activeModels)
|
||||
{
|
||||
lock (_modelSnapshotLock)
|
||||
{
|
||||
if (models is not null)
|
||||
{
|
||||
_lastModels = models;
|
||||
}
|
||||
|
||||
if (activeModels is not null)
|
||||
{
|
||||
_lastActiveModels = activeModels;
|
||||
}
|
||||
|
||||
return (
|
||||
[.. _lastModels],
|
||||
[.. _lastActiveModels]);
|
||||
}
|
||||
}
|
||||
|
||||
private static List<string> NormalizeModelNames(IEnumerable<string> models) =>
|
||||
models
|
||||
.Where(model => !string.IsNullOrWhiteSpace(model))
|
||||
.Select(model => model.Trim())
|
||||
.Distinct(StringComparer.OrdinalIgnoreCase)
|
||||
.OrderBy(model => model, StringComparer.OrdinalIgnoreCase)
|
||||
.ToList();
|
||||
|
||||
private async Task ReceiveLoopAsync(ClientWebSocket socket, CancellationToken cancellationToken)
|
||||
{
|
||||
while (socket.State == WebSocketState.Open && !cancellationToken.IsCancellationRequested)
|
||||
{
|
||||
var message = await WebSocketMessageTransport.ReceiveAsync(socket, cancellationToken);
|
||||
if (message is null)
|
||||
{
|
||||
break;
|
||||
}
|
||||
|
||||
await DispatchAsync(socket, message, cancellationToken);
|
||||
}
|
||||
}
|
||||
|
||||
private Task DispatchAsync(ClientWebSocket socket, TunnelMessage message, CancellationToken cancellationToken)
|
||||
{
|
||||
switch (message.Type)
|
||||
{
|
||||
case TunnelMessageTypes.HttpRequest:
|
||||
StartRequest(socket, message, cancellationToken);
|
||||
break;
|
||||
|
||||
case TunnelMessageTypes.HttpRequestBody:
|
||||
if (_activeRequests.TryGetValue(message.RequestId, out var requestWithBody))
|
||||
{
|
||||
requestWithBody.AddBody(message.Body ?? []);
|
||||
}
|
||||
break;
|
||||
|
||||
case TunnelMessageTypes.HttpRequestComplete:
|
||||
if (_activeRequests.TryGetValue(message.RequestId, out var completedRequest))
|
||||
{
|
||||
completedRequest.CompleteBody();
|
||||
}
|
||||
break;
|
||||
|
||||
case TunnelMessageTypes.Cancel:
|
||||
if (_activeRequests.TryRemove(message.RequestId, out var cancelledRequest))
|
||||
{
|
||||
cancelledRequest.Cancel();
|
||||
}
|
||||
break;
|
||||
|
||||
case TunnelMessageTypes.ModelCommand:
|
||||
_ = Task.Run(() => RunModelCommandAsync(socket, message, cancellationToken), cancellationToken);
|
||||
break;
|
||||
}
|
||||
|
||||
return Task.CompletedTask;
|
||||
}
|
||||
|
||||
private async Task RunModelCommandAsync(ClientWebSocket socket, TunnelMessage message, CancellationToken cancellationToken)
|
||||
{
|
||||
try
|
||||
{
|
||||
var response = await ExecuteModelCommandAsync(message, cancellationToken);
|
||||
if (response.StatusCode is >= 200 and < 300)
|
||||
{
|
||||
await RefreshModelsOnceAsync(socket, cancellationToken);
|
||||
}
|
||||
|
||||
await SendAsync(socket, response, cancellationToken);
|
||||
}
|
||||
catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
|
||||
{
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
await SendAsync(
|
||||
socket,
|
||||
new TunnelMessage
|
||||
{
|
||||
Type = TunnelMessageTypes.ModelCommandResult,
|
||||
RequestId = message.RequestId,
|
||||
Error = exception.Message
|
||||
},
|
||||
CancellationToken.None);
|
||||
}
|
||||
}
|
||||
|
||||
private async Task<TunnelMessage> ExecuteModelCommandAsync(TunnelMessage message, CancellationToken cancellationToken)
|
||||
{
|
||||
if (string.IsNullOrWhiteSpace(message.RequestId))
|
||||
{
|
||||
throw new InvalidOperationException("Model command is missing a request id.");
|
||||
}
|
||||
|
||||
if (string.IsNullOrWhiteSpace(message.Model))
|
||||
{
|
||||
throw new InvalidOperationException("Model command is missing a model name.");
|
||||
}
|
||||
|
||||
using var request = BuildModelCommandRequest(message);
|
||||
using var response = await _httpClient.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, cancellationToken);
|
||||
var body = await response.Content.ReadAsByteArrayAsync(cancellationToken);
|
||||
|
||||
if (ShouldRetryModelCommandWithEmbedding(message.Command, response, body))
|
||||
{
|
||||
using var embeddingRequest = BuildEmbeddingModelCommandRequest(_options.Upstream, message.Command, message.Model);
|
||||
using var embeddingResponse = await _httpClient.SendAsync(embeddingRequest, HttpCompletionOption.ResponseHeadersRead, cancellationToken);
|
||||
var embeddingBody = await embeddingResponse.Content.ReadAsByteArrayAsync(cancellationToken);
|
||||
|
||||
return BuildModelCommandResult(message.RequestId, embeddingResponse, embeddingBody);
|
||||
}
|
||||
|
||||
return BuildModelCommandResult(message.RequestId, response, body);
|
||||
}
|
||||
|
||||
private static TunnelMessage BuildModelCommandResult(
|
||||
string requestId,
|
||||
HttpResponseMessage response,
|
||||
byte[] body)
|
||||
{
|
||||
return new TunnelMessage
|
||||
{
|
||||
Type = TunnelMessageTypes.ModelCommandResult,
|
||||
RequestId = requestId,
|
||||
StatusCode = (int)response.StatusCode,
|
||||
ReasonPhrase = response.ReasonPhrase,
|
||||
Body = body
|
||||
};
|
||||
}
|
||||
|
||||
private static bool ShouldRetryModelCommandWithEmbedding(
|
||||
string? command,
|
||||
HttpResponseMessage response,
|
||||
byte[] body)
|
||||
{
|
||||
var normalizedCommand = NormalizeModelCommand(command);
|
||||
if (normalizedCommand is not ("load" or "unload")
|
||||
|| response.IsSuccessStatusCode)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
if (response.StatusCode == System.Net.HttpStatusCode.BadRequest)
|
||||
{
|
||||
return true;
|
||||
}
|
||||
|
||||
var responseText = body.Length > 0
|
||||
? Encoding.UTF8.GetString(body)
|
||||
: "";
|
||||
|
||||
return responseText.Contains("does not support generate", StringComparison.OrdinalIgnoreCase);
|
||||
}
|
||||
|
||||
private HttpRequestMessage BuildModelCommandRequest(TunnelMessage message)
|
||||
{
|
||||
return BuildModelCommandRequest(_options.Upstream, message.Command, message.Model);
|
||||
}
|
||||
|
||||
internal static HttpRequestMessage BuildModelCommandRequest(Uri upstream, string? command, string? modelName)
|
||||
{
|
||||
if (string.IsNullOrWhiteSpace(modelName))
|
||||
{
|
||||
throw new InvalidOperationException("Model command is missing a model name.");
|
||||
}
|
||||
|
||||
var model = modelName.Trim();
|
||||
var normalizedCommand = NormalizeModelCommand(command);
|
||||
|
||||
return normalizedCommand switch
|
||||
{
|
||||
"pull" => new HttpRequestMessage(HttpMethod.Post, new Uri(upstream, "/api/pull"))
|
||||
{
|
||||
Content = JsonContent(new { model, stream = false })
|
||||
},
|
||||
"delete" => new HttpRequestMessage(HttpMethod.Delete, new Uri(upstream, "/api/delete"))
|
||||
{
|
||||
Content = JsonContent(new { model })
|
||||
},
|
||||
"load" => new HttpRequestMessage(HttpMethod.Post, new Uri(upstream, "/api/generate"))
|
||||
{
|
||||
Content = JsonContent(new { model, stream = false, keep_alive = -1 })
|
||||
},
|
||||
"unload" => new HttpRequestMessage(HttpMethod.Post, new Uri(upstream, "/api/generate"))
|
||||
{
|
||||
Content = JsonContent(new { model, stream = false, keep_alive = 0 })
|
||||
},
|
||||
"show" => new HttpRequestMessage(HttpMethod.Post, new Uri(upstream, "/api/show"))
|
||||
{
|
||||
Content = JsonContent(new { model })
|
||||
},
|
||||
_ => throw new InvalidOperationException($"Unsupported model command '{command}'.")
|
||||
};
|
||||
}
|
||||
|
||||
internal static HttpRequestMessage BuildEmbeddingModelCommandRequest(Uri upstream, string? command, string? modelName)
|
||||
{
|
||||
if (string.IsNullOrWhiteSpace(modelName))
|
||||
{
|
||||
throw new InvalidOperationException("Model command is missing a model name.");
|
||||
}
|
||||
|
||||
var model = modelName.Trim();
|
||||
var normalizedCommand = NormalizeModelCommand(command);
|
||||
|
||||
return normalizedCommand switch
|
||||
{
|
||||
"load" => new HttpRequestMessage(HttpMethod.Post, new Uri(upstream, "/api/embed"))
|
||||
{
|
||||
Content = JsonContent(new { model, input = EmbeddingWarmupInput, keep_alive = -1 })
|
||||
},
|
||||
"unload" => new HttpRequestMessage(HttpMethod.Post, new Uri(upstream, "/api/embed"))
|
||||
{
|
||||
Content = JsonContent(new { model, input = EmbeddingWarmupInput, keep_alive = 0 })
|
||||
},
|
||||
_ => throw new InvalidOperationException($"Unsupported embedding model command '{command}'.")
|
||||
};
|
||||
}
|
||||
|
||||
private static string NormalizeModelCommand(string? command) =>
|
||||
(command ?? "").Trim().ToLowerInvariant();
|
||||
|
||||
private static StringContent JsonContent<T>(T value) =>
|
||||
new(JsonSerializer.Serialize(value, JsonOptions), Encoding.UTF8, "application/json");
|
||||
|
||||
private void StartRequest(ClientWebSocket socket, TunnelMessage message, CancellationToken cancellationToken)
|
||||
{
|
||||
var request = new UpstreamRequest(
|
||||
_options,
|
||||
_httpClient,
|
||||
message,
|
||||
(response, token) => SendAsync(socket, response, token),
|
||||
requestId => _activeRequests.TryRemove(requestId, out _),
|
||||
cancellationToken);
|
||||
|
||||
if (!_activeRequests.TryAdd(message.RequestId, request))
|
||||
{
|
||||
_ = SendAsync(
|
||||
socket,
|
||||
new TunnelMessage
|
||||
{
|
||||
Type = TunnelMessageTypes.Error,
|
||||
RequestId = message.RequestId,
|
||||
Error = "Duplicate request id."
|
||||
},
|
||||
cancellationToken);
|
||||
return;
|
||||
}
|
||||
|
||||
_ = Task.Run(request.RunAsync, cancellationToken);
|
||||
}
|
||||
|
||||
private async Task SendAsync(ClientWebSocket socket, TunnelMessage message, CancellationToken cancellationToken)
|
||||
{
|
||||
await _sendLock.WaitAsync(cancellationToken);
|
||||
try
|
||||
{
|
||||
if (socket.State == WebSocketState.Open)
|
||||
{
|
||||
await WebSocketMessageTransport.SendAsync(socket, message, cancellationToken);
|
||||
}
|
||||
}
|
||||
finally
|
||||
{
|
||||
_sendLock.Release();
|
||||
}
|
||||
}
|
||||
|
||||
private void CancelAllActiveRequests()
|
||||
{
|
||||
foreach (var pair in _activeRequests.ToArray())
|
||||
{
|
||||
if (_activeRequests.TryRemove(pair.Key, out var request))
|
||||
{
|
||||
request.Cancel();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
using Microsoft.Extensions.Hosting;
|
||||
|
||||
namespace ReverseLlama.Client;
|
||||
|
||||
internal sealed class TunnelWorker : BackgroundService
|
||||
{
|
||||
private readonly TunnelClient _client;
|
||||
private readonly IHostApplicationLifetime _lifetime;
|
||||
|
||||
public TunnelWorker(TunnelClient client, IHostApplicationLifetime lifetime)
|
||||
{
|
||||
_client = client;
|
||||
_lifetime = lifetime;
|
||||
}
|
||||
|
||||
protected override async Task ExecuteAsync(CancellationToken stoppingToken)
|
||||
{
|
||||
try
|
||||
{
|
||||
await _client.RunAsync(stoppingToken);
|
||||
}
|
||||
finally
|
||||
{
|
||||
// RunAsync only returns when cancelled or replaced by a newer client.
|
||||
// Stop gracefully (exit 0) so service recovery does not restart us
|
||||
// into a reconnect fight with the replacement.
|
||||
_lifetime.StopApplication();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,245 @@
|
||||
using System.Net.Http.Headers;
|
||||
using System.Threading.Channels;
|
||||
using ReverseLlama.Protocol;
|
||||
|
||||
namespace ReverseLlama.Client;
|
||||
|
||||
internal sealed class UpstreamRequest
|
||||
{
|
||||
private static readonly HashSet<string> HeadersToSkip = new(StringComparer.OrdinalIgnoreCase)
|
||||
{
|
||||
"Connection",
|
||||
"Content-Length",
|
||||
"Expect",
|
||||
"Host",
|
||||
"Keep-Alive",
|
||||
"Proxy-Authenticate",
|
||||
"Proxy-Authorization",
|
||||
"TE",
|
||||
"Trailer",
|
||||
"Transfer-Encoding",
|
||||
"Upgrade",
|
||||
ProtocolConstants.TokenHeader
|
||||
};
|
||||
|
||||
private readonly CancellationTokenSource _cancellationTokenSource;
|
||||
private readonly Channel<byte[]> _requestBody = Channel.CreateUnbounded<byte[]>(
|
||||
new UnboundedChannelOptions
|
||||
{
|
||||
SingleReader = true,
|
||||
SingleWriter = false
|
||||
});
|
||||
|
||||
private readonly HttpClient _httpClient;
|
||||
private readonly TunnelMessage _initialMessage;
|
||||
private readonly Action<string> _onComplete;
|
||||
private readonly ClientOptions _options;
|
||||
private readonly Func<TunnelMessage, CancellationToken, Task> _sendAsync;
|
||||
|
||||
public UpstreamRequest(
|
||||
ClientOptions options,
|
||||
HttpClient httpClient,
|
||||
TunnelMessage initialMessage,
|
||||
Func<TunnelMessage, CancellationToken, Task> sendAsync,
|
||||
Action<string> onComplete,
|
||||
CancellationToken cancellationToken)
|
||||
{
|
||||
_options = options;
|
||||
_httpClient = httpClient;
|
||||
_initialMessage = initialMessage;
|
||||
_sendAsync = sendAsync;
|
||||
_onComplete = onComplete;
|
||||
_cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
|
||||
|
||||
if (!initialMessage.HasBody)
|
||||
{
|
||||
_requestBody.Writer.TryComplete();
|
||||
}
|
||||
}
|
||||
|
||||
public void AddBody(byte[] body)
|
||||
{
|
||||
if (body.Length > 0)
|
||||
{
|
||||
_requestBody.Writer.TryWrite(body);
|
||||
}
|
||||
}
|
||||
|
||||
public void CompleteBody() =>
|
||||
_requestBody.Writer.TryComplete();
|
||||
|
||||
public void Cancel()
|
||||
{
|
||||
_requestBody.Writer.TryComplete();
|
||||
_cancellationTokenSource.Cancel();
|
||||
}
|
||||
|
||||
public async Task RunAsync()
|
||||
{
|
||||
try
|
||||
{
|
||||
using var request = BuildHttpRequest();
|
||||
using var response = await _httpClient.SendAsync(
|
||||
request,
|
||||
HttpCompletionOption.ResponseHeadersRead,
|
||||
_cancellationTokenSource.Token);
|
||||
|
||||
await SendResponseHeadersAsync(response);
|
||||
await SendResponseBodyAsync(response);
|
||||
|
||||
await _sendAsync(
|
||||
new TunnelMessage
|
||||
{
|
||||
Type = TunnelMessageTypes.HttpResponseComplete,
|
||||
RequestId = _initialMessage.RequestId
|
||||
},
|
||||
_cancellationTokenSource.Token);
|
||||
}
|
||||
catch (OperationCanceledException) when (_cancellationTokenSource.IsCancellationRequested)
|
||||
{
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
await SendErrorAsync(exception);
|
||||
}
|
||||
finally
|
||||
{
|
||||
_requestBody.Writer.TryComplete();
|
||||
_onComplete(_initialMessage.RequestId);
|
||||
_cancellationTokenSource.Dispose();
|
||||
}
|
||||
}
|
||||
|
||||
private HttpRequestMessage BuildHttpRequest()
|
||||
{
|
||||
var method = new HttpMethod(_initialMessage.Method ?? HttpMethod.Get.Method);
|
||||
var request = new HttpRequestMessage(method, BuildUpstreamUri(_options.Upstream, _initialMessage.PathAndQuery));
|
||||
|
||||
if (_initialMessage.HasBody)
|
||||
{
|
||||
request.Content = new ChannelHttpContent(_requestBody.Reader, _cancellationTokenSource.Token);
|
||||
}
|
||||
|
||||
foreach (var header in _initialMessage.Headers)
|
||||
{
|
||||
if (HeadersToSkip.Contains(header.Name))
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
if (!request.Headers.TryAddWithoutValidation(header.Name, header.Value)
|
||||
&& request.Content is not null)
|
||||
{
|
||||
request.Content.Headers.TryAddWithoutValidation(header.Name, header.Value);
|
||||
}
|
||||
}
|
||||
|
||||
return request;
|
||||
}
|
||||
|
||||
internal static Uri BuildUpstreamUri(Uri upstream, string? pathAndQuery)
|
||||
{
|
||||
if (string.IsNullOrWhiteSpace(pathAndQuery))
|
||||
{
|
||||
return upstream;
|
||||
}
|
||||
|
||||
if (!IsOriginPathAndQuery(pathAndQuery))
|
||||
{
|
||||
throw new InvalidOperationException("Tunnel request path must be an origin-form path.");
|
||||
}
|
||||
|
||||
var uri = new Uri(upstream, pathAndQuery);
|
||||
if (!HasSameOrigin(uri, upstream))
|
||||
{
|
||||
throw new InvalidOperationException("Tunnel request path resolved outside the configured upstream origin.");
|
||||
}
|
||||
|
||||
return uri;
|
||||
}
|
||||
|
||||
private static bool IsOriginPathAndQuery(string pathAndQuery) =>
|
||||
pathAndQuery.StartsWith("/", StringComparison.Ordinal)
|
||||
&& !pathAndQuery.StartsWith("//", StringComparison.Ordinal)
|
||||
&& !pathAndQuery.Contains('\\');
|
||||
|
||||
private static bool HasSameOrigin(Uri uri, Uri upstream) =>
|
||||
string.Equals(uri.Scheme, upstream.Scheme, StringComparison.OrdinalIgnoreCase)
|
||||
&& string.Equals(uri.IdnHost, upstream.IdnHost, StringComparison.OrdinalIgnoreCase)
|
||||
&& uri.Port == upstream.Port;
|
||||
|
||||
private async Task SendResponseHeadersAsync(HttpResponseMessage response)
|
||||
{
|
||||
await _sendAsync(
|
||||
new TunnelMessage
|
||||
{
|
||||
Type = TunnelMessageTypes.HttpResponseHeaders,
|
||||
RequestId = _initialMessage.RequestId,
|
||||
StatusCode = (int)response.StatusCode,
|
||||
ReasonPhrase = response.ReasonPhrase,
|
||||
Headers = CollectResponseHeaders(response)
|
||||
},
|
||||
_cancellationTokenSource.Token);
|
||||
}
|
||||
|
||||
private async Task SendResponseBodyAsync(HttpResponseMessage response)
|
||||
{
|
||||
await using var stream = await response.Content.ReadAsStreamAsync(_cancellationTokenSource.Token);
|
||||
var buffer = new byte[_options.ChunkSize];
|
||||
|
||||
while (true)
|
||||
{
|
||||
var bytesRead = await stream.ReadAsync(buffer, _cancellationTokenSource.Token);
|
||||
if (bytesRead == 0)
|
||||
{
|
||||
break;
|
||||
}
|
||||
|
||||
await _sendAsync(
|
||||
new TunnelMessage
|
||||
{
|
||||
Type = TunnelMessageTypes.HttpResponseBody,
|
||||
RequestId = _initialMessage.RequestId,
|
||||
Body = buffer.AsSpan(0, bytesRead).ToArray()
|
||||
},
|
||||
_cancellationTokenSource.Token);
|
||||
}
|
||||
}
|
||||
|
||||
private async Task SendErrorAsync(Exception exception)
|
||||
{
|
||||
try
|
||||
{
|
||||
await _sendAsync(
|
||||
new TunnelMessage
|
||||
{
|
||||
Type = TunnelMessageTypes.Error,
|
||||
RequestId = _initialMessage.RequestId,
|
||||
Error = exception.Message
|
||||
},
|
||||
CancellationToken.None);
|
||||
}
|
||||
catch
|
||||
{
|
||||
}
|
||||
}
|
||||
|
||||
private static List<HeaderPair> CollectResponseHeaders(HttpResponseMessage response)
|
||||
{
|
||||
var headers = new List<HeaderPair>();
|
||||
AddHeaders(headers, response.Headers);
|
||||
AddHeaders(headers, response.Content.Headers);
|
||||
return headers;
|
||||
}
|
||||
|
||||
private static void AddHeaders(List<HeaderPair> target, HttpHeaders headers)
|
||||
{
|
||||
foreach (var header in headers)
|
||||
{
|
||||
foreach (var value in header.Value)
|
||||
{
|
||||
target.Add(new HeaderPair(header.Key, value));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user