429 lines
15 KiB
C#
429 lines
15 KiB
C#
// File: OperationCenter.cs
|
|
using System;
|
|
using System.Collections.Generic;
|
|
using System.Linq;
|
|
using System.Text.Json.Nodes;
|
|
|
|
namespace AI_C2_Server
|
|
{
|
|
public class OperationCenter
|
|
{
|
|
public static OperationCenter Instance { get; } = new OperationCenter();
|
|
|
|
private readonly Dictionary<string, Queue<JsonObject>> commandQueues = new();
|
|
private readonly Dictionary<string, JsonObject> results = new();
|
|
private readonly Dictionary<string, List<string>> pendingTaskIds = new();
|
|
private readonly Dictionary<string, JsonObject> instances = new();
|
|
private readonly Dictionary<string, JsonArray> chatHistory = new();
|
|
private readonly Dictionary<string, JsonObject> tasks = new();
|
|
|
|
// Tracks instances that have been shut down so they aren't recreated by lingering heartbeats
|
|
private readonly HashSet<string> blacklistedInstances = new();
|
|
|
|
private readonly object queueLock = new object();
|
|
|
|
// OpenAI Configuration
|
|
public bool UseOpenAIMode { get; set; } = false;
|
|
public string OpenAIApiUrl { get; set; } = "http://192.168.12.181:8080/v1/chat/completions";
|
|
public string OpenAIApiKey { get; set; } = "";
|
|
public string OpenAIInstanceId { get; private set; }
|
|
public string OpenAIModelName { get; set; } = "";
|
|
public string OpenAISystemInstructions { get; set; } = "";
|
|
|
|
private OperationCenter() { }
|
|
|
|
public void ClearAll()
|
|
{
|
|
lock (queueLock)
|
|
{
|
|
commandQueues.Clear();
|
|
results.Clear();
|
|
pendingTaskIds.Clear();
|
|
instances.Clear();
|
|
chatHistory.Clear();
|
|
tasks.Clear();
|
|
blacklistedInstances.Clear();
|
|
UseOpenAIMode = false;
|
|
OpenAIInstanceId = null;
|
|
OpenAIModelName = "";
|
|
OpenAISystemInstructions = "";
|
|
}
|
|
}
|
|
|
|
public bool IsOpenAIInstance(string instanceId)
|
|
{
|
|
return !string.IsNullOrEmpty(OpenAIInstanceId) && instanceId == OpenAIInstanceId;
|
|
}
|
|
|
|
public string EnableOpenAIMode(string url, string key)
|
|
{
|
|
lock (queueLock)
|
|
{
|
|
UseOpenAIMode = true;
|
|
OpenAIApiUrl = url;
|
|
OpenAIApiKey = key;
|
|
OpenAIModelName = ""; // Reset model name so it fetches on next run
|
|
|
|
if (string.IsNullOrEmpty(OpenAIInstanceId))
|
|
{
|
|
OpenAIInstanceId = $"inst_openai_{DateTimeOffset.UtcNow.ToUnixTimeSeconds()}";
|
|
}
|
|
|
|
instances[OpenAIInstanceId] = new JsonObject
|
|
{
|
|
["last_seen"] = DateTime.Now.ToString("o"),
|
|
["url"] = url,
|
|
["turn_count"] = 0,
|
|
["is_busy"] = false,
|
|
["platform"] = "OpenAI API"
|
|
};
|
|
|
|
if (!chatHistory.ContainsKey(OpenAIInstanceId))
|
|
{
|
|
chatHistory[OpenAIInstanceId] = new JsonArray();
|
|
}
|
|
|
|
Logger.Log($"[OP-CENTER] OpenAI API instance registered: {OpenAIInstanceId}");
|
|
return OpenAIInstanceId;
|
|
}
|
|
}
|
|
|
|
public void DisableOpenAIMode()
|
|
{
|
|
lock (queueLock)
|
|
{
|
|
UseOpenAIMode = false;
|
|
OpenAIModelName = "";
|
|
if (!string.IsNullOrEmpty(OpenAIInstanceId))
|
|
{
|
|
instances.Remove(OpenAIInstanceId);
|
|
Logger.Log($"[OP-CENTER] OpenAI API instance removed: {OpenAIInstanceId}");
|
|
OpenAIInstanceId = null;
|
|
}
|
|
}
|
|
}
|
|
|
|
public void UpdateOpenAIUrl(string url)
|
|
{
|
|
lock (queueLock)
|
|
{
|
|
OpenAIModelName = ""; // Reset model name so it fetches again for the new URL
|
|
if (!string.IsNullOrEmpty(OpenAIInstanceId) && instances.ContainsKey(OpenAIInstanceId))
|
|
{
|
|
instances[OpenAIInstanceId]["url"] = url;
|
|
}
|
|
}
|
|
}
|
|
|
|
private void EnsureInstance(string instanceId)
|
|
{
|
|
bool isNew = false;
|
|
if (!commandQueues.ContainsKey(instanceId))
|
|
{
|
|
commandQueues[instanceId] = new Queue<JsonObject>();
|
|
pendingTaskIds[instanceId] = new List<string>();
|
|
isNew = true;
|
|
}
|
|
|
|
if (!blacklistedInstances.Contains(instanceId) && !instances.ContainsKey(instanceId))
|
|
{
|
|
instances[instanceId] = new JsonObject
|
|
{
|
|
["last_seen"] = DateTime.Now.ToString("o"),
|
|
["url"] = "",
|
|
["turn_count"] = 0,
|
|
["is_busy"] = false,
|
|
["platform"] = "Unknown"
|
|
};
|
|
|
|
if (isNew)
|
|
{
|
|
Logger.Log($"[OP-CENTER] New instance registered: {instanceId}");
|
|
}
|
|
}
|
|
|
|
if (!chatHistory.ContainsKey(instanceId))
|
|
{
|
|
chatHistory[instanceId] = new JsonArray();
|
|
}
|
|
}
|
|
|
|
public void UpdateInstanceState(string instanceId, JsonObject stateData)
|
|
{
|
|
lock (queueLock)
|
|
{
|
|
if (blacklistedInstances.Contains(instanceId)) return;
|
|
|
|
EnsureInstance(instanceId);
|
|
|
|
if (instances.ContainsKey(instanceId))
|
|
{
|
|
instances[instanceId]["last_seen"] = DateTime.Now.ToString("o");
|
|
|
|
if (stateData != null)
|
|
{
|
|
if (stateData.ContainsKey("url"))
|
|
instances[instanceId]["url"] = stateData["url"]?.ToString();
|
|
if (stateData.ContainsKey("turn_count"))
|
|
instances[instanceId]["turn_count"] = stateData["turn_count"]?.GetValue<int>() ?? 0;
|
|
if (stateData.ContainsKey("is_busy"))
|
|
instances[instanceId]["is_busy"] = stateData["is_busy"]?.GetValue<bool>() ?? false;
|
|
if (stateData.ContainsKey("platform"))
|
|
instances[instanceId]["platform"] = stateData["platform"]?.ToString();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
public JsonObject GetActiveInstances()
|
|
{
|
|
lock (queueLock)
|
|
{
|
|
var result = new JsonObject();
|
|
foreach (var kvp in instances)
|
|
{
|
|
result[kvp.Key] = kvp.Value.DeepClone();
|
|
}
|
|
return result;
|
|
}
|
|
}
|
|
|
|
public string GetDefaultInstanceId()
|
|
{
|
|
lock (queueLock)
|
|
{
|
|
if (instances.Count == 0) return "default";
|
|
return instances.OrderByDescending(x => DateTime.Parse(x.Value["last_seen"].ToString())).First().Key;
|
|
}
|
|
}
|
|
|
|
public string GetInstancePlatform(string instanceId)
|
|
{
|
|
lock (queueLock)
|
|
{
|
|
if (instances.TryGetValue(instanceId, out var data))
|
|
{
|
|
return data?["platform"]?.ToString() ?? "Unknown";
|
|
}
|
|
return "Unknown";
|
|
}
|
|
}
|
|
|
|
public string QueueCommand(JsonObject commandData, string instanceId = null)
|
|
{
|
|
lock (queueLock)
|
|
{
|
|
if (string.IsNullOrEmpty(instanceId))
|
|
{
|
|
instanceId = GetDefaultInstanceId();
|
|
}
|
|
|
|
EnsureInstance(instanceId);
|
|
|
|
string taskId = $"CMD-{instanceId.Substring(0, Math.Min(8, instanceId.Length))}-{DateTimeOffset.UtcNow.ToUnixTimeSeconds()}";
|
|
|
|
var command = new JsonObject
|
|
{
|
|
["id"] = taskId,
|
|
["command"] = commandData.DeepClone(),
|
|
["timestamp"] = DateTime.Now.ToString("o"),
|
|
["instance_id"] = instanceId
|
|
};
|
|
|
|
tasks[taskId] = commandData.DeepClone().AsObject();
|
|
|
|
string cmdType = commandData["type"]?.ToString();
|
|
if (cmdType == "GENERATE")
|
|
{
|
|
chatHistory[instanceId].Add(new JsonObject
|
|
{
|
|
["role"] = "User",
|
|
["text"] = commandData["payload"]?.ToString() ?? ""
|
|
});
|
|
}
|
|
|
|
commandQueues[instanceId].Enqueue(command);
|
|
pendingTaskIds[instanceId].Add(taskId);
|
|
|
|
return taskId;
|
|
}
|
|
}
|
|
|
|
public JsonObject GetPendingCommand(string instanceId)
|
|
{
|
|
lock (queueLock)
|
|
{
|
|
if (!commandQueues.ContainsKey(instanceId)) return null;
|
|
|
|
if (commandQueues[instanceId].Count > 0)
|
|
{
|
|
if (instances.ContainsKey(instanceId))
|
|
{
|
|
instances[instanceId]["is_busy"] = true;
|
|
}
|
|
return commandQueues[instanceId].Dequeue();
|
|
}
|
|
return null;
|
|
}
|
|
}
|
|
|
|
public bool LogResult(string instanceId, JsonObject resultData)
|
|
{
|
|
lock (queueLock)
|
|
{
|
|
if (instances.ContainsKey(instanceId))
|
|
{
|
|
instances[instanceId]["is_busy"] = false;
|
|
}
|
|
|
|
if (pendingTaskIds.ContainsKey(instanceId) && pendingTaskIds[instanceId].Count > 0)
|
|
{
|
|
string matchedTaskId = pendingTaskIds[instanceId][0];
|
|
pendingTaskIds[instanceId].RemoveAt(0);
|
|
|
|
if (tasks.TryGetValue(matchedTaskId, out var cmdData))
|
|
{
|
|
string cmdType = cmdData["type"]?.ToString();
|
|
string output = resultData["output"]?.ToString() ?? "";
|
|
|
|
if (cmdType == "GET_CHAT_HISTORY")
|
|
{
|
|
try { chatHistory[instanceId] = JsonNode.Parse(output).AsArray(); } catch { }
|
|
}
|
|
else if ((cmdType == "GENERATE" || cmdType == "INJECT_SHARD") && !output.StartsWith("Error"))
|
|
{
|
|
chatHistory[instanceId].Add(new JsonObject { ["role"] = "Model", ["text"] = output });
|
|
}
|
|
else if (cmdType == "NEW_CHAT")
|
|
{
|
|
chatHistory[instanceId].Clear();
|
|
}
|
|
else if (cmdType == "DELETE_TURN" && !output.StartsWith("Error"))
|
|
{
|
|
try
|
|
{
|
|
int idx = cmdData["payload"]?.GetValue<int>() ?? -1;
|
|
if (idx < 0) idx = chatHistory[instanceId].Count + idx;
|
|
if (idx >= 0 && idx < chatHistory[instanceId].Count)
|
|
{
|
|
chatHistory[instanceId].RemoveAt(idx);
|
|
}
|
|
}
|
|
catch { }
|
|
}
|
|
}
|
|
|
|
results[matchedTaskId] = new JsonObject
|
|
{
|
|
["received_at"] = DateTime.Now.ToString("o"),
|
|
["data"] = resultData.DeepClone(),
|
|
["instance_id"] = instanceId
|
|
};
|
|
|
|
Logger.Log($"[OP-CENTER] Task {matchedTaskId} Completed on Instance {instanceId}.");
|
|
return true;
|
|
}
|
|
else
|
|
{
|
|
Logger.Log($"[OP-CENTER] Warning: Result received from {instanceId} but no tasks were pending.");
|
|
return false;
|
|
}
|
|
}
|
|
}
|
|
|
|
public JsonObject GetResult(string taskId)
|
|
{
|
|
lock (queueLock)
|
|
{
|
|
if (results.TryGetValue(taskId, out var res))
|
|
{
|
|
return res.DeepClone().AsObject();
|
|
}
|
|
return null;
|
|
}
|
|
}
|
|
|
|
public void CancelPendingTask(string instanceId, string taskId)
|
|
{
|
|
lock (queueLock)
|
|
{
|
|
if (pendingTaskIds.ContainsKey(instanceId))
|
|
{
|
|
if (pendingTaskIds[instanceId].Remove(taskId))
|
|
{
|
|
Logger.Log($"[OP-CENTER] Task {taskId} timed out and was removed from pending queue for {instanceId}.");
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
public void ShutdownInstance(string instanceId)
|
|
{
|
|
lock (queueLock)
|
|
{
|
|
if (IsOpenAIInstance(instanceId))
|
|
{
|
|
DisableOpenAIMode();
|
|
return;
|
|
}
|
|
|
|
// Queue the shutdown command so the harness receives it on its next poll
|
|
QueueCommand(new JsonObject { ["type"] = "SHUTDOWN" }, instanceId);
|
|
|
|
// Blacklist and remove from active instances so it disappears from the UI
|
|
// and cannot be recreated by lingering heartbeats
|
|
blacklistedInstances.Add(instanceId);
|
|
instances.Remove(instanceId);
|
|
}
|
|
}
|
|
|
|
public void SetInstanceBusy(string instanceId, bool isBusy)
|
|
{
|
|
lock (queueLock)
|
|
{
|
|
if (instances.ContainsKey(instanceId))
|
|
{
|
|
instances[instanceId]["is_busy"] = isBusy;
|
|
instances[instanceId]["last_seen"] = DateTime.Now.ToString("o");
|
|
}
|
|
}
|
|
}
|
|
|
|
public void AddChatHistory(string instanceId, string role, string text)
|
|
{
|
|
lock (queueLock)
|
|
{
|
|
if (!chatHistory.ContainsKey(instanceId)) chatHistory[instanceId] = new JsonArray();
|
|
chatHistory[instanceId].Add(new JsonObject { ["role"] = role, ["text"] = text });
|
|
}
|
|
}
|
|
|
|
public void ClearChatHistory(string instanceId)
|
|
{
|
|
lock (queueLock)
|
|
{
|
|
if (chatHistory.ContainsKey(instanceId)) chatHistory[instanceId].Clear();
|
|
}
|
|
}
|
|
|
|
public bool HasChatHistory(string instanceId)
|
|
{
|
|
lock (queueLock)
|
|
{
|
|
return chatHistory.ContainsKey(instanceId) && chatHistory[instanceId].Count > 0;
|
|
}
|
|
}
|
|
|
|
public JsonArray GetChatHistory(string instanceId)
|
|
{
|
|
lock (queueLock)
|
|
{
|
|
if (chatHistory.TryGetValue(instanceId, out var history))
|
|
{
|
|
return history.DeepClone().AsArray();
|
|
}
|
|
return new JsonArray();
|
|
}
|
|
}
|
|
}
|
|
} |