wehub-resource-sync
534 lines
23 KiB
C#
534 lines
23 KiB
C#
using System;
|
|
using System.Collections.Generic;
|
|
using System.IO;
|
|
using System.Threading;
|
|
using System.Threading.Tasks;
|
|
using MCPForUnity.Editor.Helpers;
|
|
using MCPForUnity.Editor.Security;
|
|
using MCPForUnity.Editor.Services.AssetGen.Http;
|
|
using MCPForUnity.Editor.Services.AssetGen.Import;
|
|
using MCPForUnity.Editor.Services.AssetGen.Providers;
|
|
using Newtonsoft.Json;
|
|
using UnityEditor;
|
|
using UnityEngine;
|
|
|
|
namespace MCPForUnity.Editor.Services.AssetGen
|
|
{
|
|
public enum AssetGenJobState { Queued, Running, Importing, Done, Failed, Canceled }
|
|
|
|
/// <summary>
|
|
/// Snapshot of a generation/import job. Persisted to SessionState so a `status` query
|
|
/// still works after an unrelated domain reload. NEVER carries a key or secret.
|
|
/// </summary>
|
|
public sealed class AssetGenJob
|
|
{
|
|
public string JobId;
|
|
public string Kind; // model | image | marketplace
|
|
public string Provider;
|
|
public string Action;
|
|
public AssetGenJobState State;
|
|
public float Progress;
|
|
public string Format;
|
|
public float TargetSize = 1f;
|
|
public string AnimationType; // FBX/OBJ rig mode: null/none | generic | humanoid | legacy
|
|
public string AssetPath;
|
|
public string AssetGuid;
|
|
public string Error;
|
|
}
|
|
|
|
/// <summary>
|
|
/// Drives asset-generation jobs on the Unity main thread via EditorApplication.update. Each
|
|
/// job runs a generic submit → poll → (download | inline) → import state machine; the model
|
|
/// and image paths supply their own submit/poll/import delegates. Because UnityWebRequest
|
|
/// completes on the main thread, polling Task.IsCompleted from the update loop is
|
|
/// main-thread-safe — we never block or use threadpool waits. The provider key is read once at
|
|
/// submit time, captured only inside the submit/poll closures — never stored on the job,
|
|
/// persisted, or logged.
|
|
/// </summary>
|
|
[InitializeOnLoad]
|
|
public static class AssetGenJobManager
|
|
{
|
|
private const string JobKeyPrefix = "MCPForUnity.AssetGen.Job.";
|
|
private const string JobIndexKey = "MCPForUnity.AssetGen.JobIndex";
|
|
|
|
// Test seams (overridable; defaults are the production implementations).
|
|
internal static IHttpTransport TransportOverrideForTests;
|
|
internal static Func<AssetGenJob, string, AssetGenJob> ImportOverrideForTests;
|
|
internal static double PollIntervalSeconds = 3.0;
|
|
internal static double TimeoutSeconds = 600.0;
|
|
|
|
private static readonly Dictionary<string, AssetGenJob> Jobs = new();
|
|
private static readonly Dictionary<string, Runner> Runners = new();
|
|
private static readonly List<string> _tickIds = new();
|
|
private static bool _ticking;
|
|
|
|
static AssetGenJobManager()
|
|
{
|
|
try
|
|
{
|
|
string index = SessionState.GetString(JobIndexKey, string.Empty);
|
|
if (string.IsNullOrEmpty(index)) return;
|
|
foreach (string id in index.Split(new[] { ',' }, StringSplitOptions.RemoveEmptyEntries))
|
|
{
|
|
string json = SessionState.GetString(JobKeyPrefix + id, string.Empty);
|
|
if (string.IsNullOrEmpty(json)) continue;
|
|
var job = JsonConvert.DeserializeObject<AssetGenJob>(json);
|
|
if (job == null) continue;
|
|
if (!IsTerminal(job.State))
|
|
{
|
|
job.State = AssetGenJobState.Failed;
|
|
job.Error = "Interrupted by an editor reload; please retry.";
|
|
Persist(job);
|
|
}
|
|
Jobs[id] = job;
|
|
}
|
|
}
|
|
catch { /* recovery is best-effort */ }
|
|
}
|
|
|
|
public static AssetGenJob StartModelGeneration(ModelGenRequest req)
|
|
{
|
|
if (req == null) throw new ArgumentNullException(nameof(req));
|
|
string provider = string.IsNullOrEmpty(req.Provider) ? "tripo" : req.Provider;
|
|
IModelProviderAdapter adapter = AssetGenProviders.Model(provider); // throws NotSupportedException if unimplemented
|
|
|
|
var job = NewJob("model", provider, "generate");
|
|
job.Format = string.IsNullOrEmpty(req.Format) ? "glb" : req.Format;
|
|
job.TargetSize = req.TargetSize <= 0 ? 1f : req.TargetSize;
|
|
|
|
if (!TryResolveKey(provider, job, out string apiKey)) return job;
|
|
|
|
IHttpTransport transport = TransportOverrideForTests ?? new UnityWebRequestTransport();
|
|
var runner = new Runner
|
|
{
|
|
Job = job,
|
|
SubmitFn = ct => adapter.SubmitAsync(req, apiKey, transport, ct),
|
|
PollFn = (pid, ct) => adapter.PollAsync(pid, apiKey, transport, ct),
|
|
ImportFn = ImportOverrideForTests ?? ModelImportPipeline.ImportInto,
|
|
Transport = transport,
|
|
OutputFolder = req.OutputFolder,
|
|
Ext = (job.Format ?? "glb").TrimStart('.'),
|
|
Name = NameFrom(req.Name, req.Prompt, job.JobId),
|
|
Subfolder = "Models",
|
|
};
|
|
Register(job, runner);
|
|
return job;
|
|
}
|
|
|
|
public static AssetGenJob StartImageGeneration(ImageGenRequest req)
|
|
{
|
|
if (req == null) throw new ArgumentNullException(nameof(req));
|
|
string provider = string.IsNullOrEmpty(req.Provider) ? "fal" : req.Provider;
|
|
IImageProviderAdapter adapter = AssetGenProviders.Image(provider); // throws NotSupportedException if unimplemented
|
|
|
|
var job = NewJob("image", provider, "generate");
|
|
job.Format = "png";
|
|
|
|
if (!TryResolveKey(provider, job, out string apiKey)) return job;
|
|
|
|
IHttpTransport transport = TransportOverrideForTests ?? new UnityWebRequestTransport();
|
|
bool asSprite = req.AsSprite;
|
|
bool transparent = req.Transparent;
|
|
var runner = new Runner
|
|
{
|
|
Job = job,
|
|
SubmitFn = ct => adapter.SubmitAsync(req, apiKey, transport, ct),
|
|
PollFn = (pid, ct) => adapter.PollAsync(pid, apiKey, transport, ct),
|
|
ImportFn = ImportOverrideForTests ?? ((j, path) => ImageImportPipeline.ImportInto(j, path, asSprite, transparent, isColor: true)),
|
|
Transport = transport,
|
|
OutputFolder = req.OutputFolder,
|
|
Ext = "png",
|
|
Name = NameFrom(req.Name, req.Prompt, job.JobId),
|
|
Subfolder = "Images",
|
|
};
|
|
Register(job, runner);
|
|
return job;
|
|
}
|
|
|
|
public static AssetGenJob StartMarketplaceImport(string uid, float targetSize, string name, string outputFolder)
|
|
{
|
|
if (string.IsNullOrEmpty(uid)) throw new ArgumentException("uid required");
|
|
var adapter = AssetGenProviders.Marketplace("sketchfab"); // throws NotSupported if unimplemented
|
|
var job = NewJob("marketplace", "sketchfab", "import");
|
|
job.TargetSize = targetSize <= 0 ? 1f : targetSize;
|
|
if (!TryResolveKey("sketchfab", job, out string apiKey)) return job;
|
|
var transport = TransportOverrideForTests ?? new UnityWebRequestTransport();
|
|
var runner = new Runner
|
|
{
|
|
Job = job,
|
|
SubmitFn = ct => adapter.ResolveDownloadUrlAsync(uid, apiKey, transport, ct), // returns the zip/gltf URL as providerJobId
|
|
PollFn = (pid, ct) => Task.FromResult(new ProviderPollResult { State = ProviderPollState.Succeeded, Progress = 1f, DownloadUrl = pid, ResultExt = "zip" }),
|
|
ImportFn = ImportOverrideForTests ?? ModelImportPipeline.ImportInto,
|
|
Transport = transport,
|
|
OutputFolder = outputFolder,
|
|
Ext = "zip",
|
|
Name = NameFrom(name, uid, job.JobId),
|
|
Subfolder = "Sketchfab",
|
|
};
|
|
Register(job, runner);
|
|
return job;
|
|
}
|
|
|
|
public static AssetGenJob GetJob(string jobId)
|
|
=> string.IsNullOrEmpty(jobId) ? null : (Jobs.TryGetValue(jobId, out var j) ? j : null);
|
|
|
|
/// <summary>Most-recent-first snapshot of known jobs (for the GUI readout). Never contains keys.</summary>
|
|
public static IReadOnlyList<AssetGenJob> RecentJobs(int max = 20)
|
|
{
|
|
var all = new List<AssetGenJob>(Jobs.Values);
|
|
int start = Math.Max(0, all.Count - max);
|
|
var slice = all.GetRange(start, all.Count - start);
|
|
slice.Reverse();
|
|
return slice;
|
|
}
|
|
|
|
public static bool Cancel(string jobId)
|
|
{
|
|
if (string.IsNullOrEmpty(jobId)) return false;
|
|
if (Runners.TryGetValue(jobId, out var r))
|
|
{
|
|
r.Canceled = true;
|
|
try { r.Cts.Cancel(); } catch { }
|
|
return true;
|
|
}
|
|
if (Jobs.TryGetValue(jobId, out var job) && job.State == AssetGenJobState.Queued)
|
|
{
|
|
job.State = AssetGenJobState.Canceled;
|
|
Persist(job);
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
// ---------- runner ----------
|
|
|
|
private enum RunnerPhase { Submit, AwaitSubmit, Poll, AwaitPoll, Download, AwaitDownload, Import }
|
|
|
|
private sealed class Runner
|
|
{
|
|
public AssetGenJob Job;
|
|
public Func<CancellationToken, Task<string>> SubmitFn;
|
|
public Func<string, CancellationToken, Task<ProviderPollResult>> PollFn;
|
|
public Func<AssetGenJob, string, AssetGenJob> ImportFn;
|
|
public IHttpTransport Transport;
|
|
public string OutputFolder;
|
|
public string Ext;
|
|
public string OverrideExt;
|
|
public string Name;
|
|
public string Subfolder;
|
|
|
|
public CancellationTokenSource Cts = new();
|
|
public RunnerPhase Phase = RunnerPhase.Submit;
|
|
public double StartedAt;
|
|
public double NextPollAt;
|
|
public string ProviderJobId;
|
|
public string DownloadUrl;
|
|
public string LocalPath;
|
|
public Task<string> SubmitTask;
|
|
public Task<ProviderPollResult> PollTask;
|
|
public Task<HttpResult> DownloadTask;
|
|
public bool Canceled;
|
|
}
|
|
|
|
private static bool TryResolveKey(string provider, AssetGenJob job, out string apiKey)
|
|
{
|
|
if (!SecureKeyStore.Current.TryGet(provider, out apiKey) || string.IsNullOrEmpty(apiKey))
|
|
{
|
|
job.State = AssetGenJobState.Failed;
|
|
job.Error = AssetGenProviders.MissingKeyMessage(provider);
|
|
Jobs[job.JobId] = job;
|
|
Persist(job);
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
private static void Register(AssetGenJob job, Runner runner)
|
|
{
|
|
runner.StartedAt = Now();
|
|
Jobs[job.JobId] = job;
|
|
Runners[job.JobId] = runner;
|
|
Persist(job);
|
|
EnsureTicking();
|
|
}
|
|
|
|
private static void EnsureTicking()
|
|
{
|
|
if (_ticking) return;
|
|
EditorApplication.update += Tick;
|
|
_ticking = true;
|
|
}
|
|
|
|
private static void Tick()
|
|
{
|
|
if (Runners.Count == 0)
|
|
{
|
|
EditorApplication.update -= Tick;
|
|
_ticking = false;
|
|
return;
|
|
}
|
|
// Snapshot keys into a reused buffer so Advance can mutate Runners mid-iteration
|
|
// without churning the GC on every editor-update frame.
|
|
_tickIds.Clear();
|
|
_tickIds.AddRange(Runners.Keys);
|
|
foreach (string id in _tickIds)
|
|
{
|
|
if (Runners.TryGetValue(id, out var r)) Advance(r);
|
|
}
|
|
}
|
|
|
|
/// <summary>Advance one job one step. Returns true when the job is terminal.</summary>
|
|
internal static bool TryAdvanceForTests(string jobId)
|
|
{
|
|
if (Runners.TryGetValue(jobId, out var r))
|
|
{
|
|
Advance(r);
|
|
return IsTerminal(r.Job.State);
|
|
}
|
|
return Jobs.TryGetValue(jobId, out var j) && IsTerminal(j.State);
|
|
}
|
|
|
|
private static void Advance(Runner r)
|
|
{
|
|
if (IsTerminal(r.Job.State)) { Finalize(r); return; }
|
|
if (r.Canceled) { r.Job.State = AssetGenJobState.Canceled; Persist(r.Job); Finalize(r); return; }
|
|
if (Now() - r.StartedAt > TimeoutSeconds) { Fail(r, $"Timed out after {TimeoutSeconds:0}s."); return; }
|
|
|
|
try
|
|
{
|
|
switch (r.Phase)
|
|
{
|
|
case RunnerPhase.Submit:
|
|
r.Job.State = AssetGenJobState.Running;
|
|
Persist(r.Job);
|
|
r.SubmitTask = r.SubmitFn(r.Cts.Token);
|
|
r.Phase = RunnerPhase.AwaitSubmit;
|
|
break;
|
|
|
|
case RunnerPhase.AwaitSubmit:
|
|
if (!r.SubmitTask.IsCompleted) break;
|
|
if (Faulted(r.SubmitTask, out string subErr)) { Fail(r, subErr); break; }
|
|
r.ProviderJobId = r.SubmitTask.Result;
|
|
if (string.IsNullOrEmpty(r.ProviderJobId)) { Fail(r, "Provider returned no job id."); break; }
|
|
r.NextPollAt = Now();
|
|
r.Phase = RunnerPhase.Poll;
|
|
break;
|
|
|
|
case RunnerPhase.Poll:
|
|
if (Now() < r.NextPollAt) break;
|
|
r.PollTask = r.PollFn(r.ProviderJobId, r.Cts.Token);
|
|
r.Phase = RunnerPhase.AwaitPoll;
|
|
break;
|
|
|
|
case RunnerPhase.AwaitPoll:
|
|
if (!r.PollTask.IsCompleted) break;
|
|
if (Faulted(r.PollTask, out string pollErr)) { Fail(r, pollErr); break; }
|
|
ProviderPollResult pr = r.PollTask.Result;
|
|
r.Job.Progress = Mathf.Clamp01(pr.Progress);
|
|
Persist(r.Job);
|
|
if (pr.State == ProviderPollState.Succeeded)
|
|
{
|
|
r.OverrideExt = pr.ResultExt;
|
|
if (pr.InlineData != null && pr.InlineData.Length > 0)
|
|
{
|
|
r.LocalPath = WriteFile(r, pr.InlineData);
|
|
r.Job.State = AssetGenJobState.Importing;
|
|
Persist(r.Job);
|
|
r.Phase = RunnerPhase.Import;
|
|
}
|
|
else if (!string.IsNullOrEmpty(pr.DownloadUrl))
|
|
{
|
|
r.DownloadUrl = pr.DownloadUrl;
|
|
r.Phase = RunnerPhase.Download;
|
|
}
|
|
else
|
|
{
|
|
Fail(r, "Provider succeeded but returned no result data.");
|
|
}
|
|
}
|
|
else if (pr.State == ProviderPollState.Failed)
|
|
{
|
|
Fail(r, string.IsNullOrEmpty(pr.Error) ? "Provider reported failure." : pr.Error);
|
|
}
|
|
else
|
|
{
|
|
r.NextPollAt = Now() + PollIntervalSeconds;
|
|
r.Phase = RunnerPhase.Poll;
|
|
}
|
|
break;
|
|
|
|
case RunnerPhase.Download:
|
|
// The download URL comes from an untrusted provider response. Only fetch
|
|
// http(s) — refuse file://, ftp://, etc. so a malicious response can't read
|
|
// a local file into the project or hit an internal host.
|
|
if (!IsAllowedDownloadUrl(r.DownloadUrl))
|
|
{
|
|
Fail(r, "Refusing to fetch a non-http(s) download URL returned by the provider.");
|
|
break;
|
|
}
|
|
r.DownloadTask = r.Transport.SendAsync(
|
|
new HttpRequestSpec { Method = "GET", Url = r.DownloadUrl }, r.Cts.Token);
|
|
r.Phase = RunnerPhase.AwaitDownload;
|
|
break;
|
|
|
|
case RunnerPhase.AwaitDownload:
|
|
if (!r.DownloadTask.IsCompleted) break;
|
|
if (Faulted(r.DownloadTask, out string dlErr)) { Fail(r, dlErr); break; }
|
|
HttpResult res = r.DownloadTask.Result;
|
|
if (res == null || !res.IsSuccess || res.Body == null || res.Body.Length == 0)
|
|
{
|
|
Fail(r, $"Download failed (HTTP {res?.Status}).");
|
|
break;
|
|
}
|
|
r.LocalPath = WriteFile(r, res.Body);
|
|
r.Job.State = AssetGenJobState.Importing;
|
|
Persist(r.Job);
|
|
r.Phase = RunnerPhase.Import;
|
|
break;
|
|
|
|
case RunnerPhase.Import:
|
|
// The result file was just written via File.WriteAllBytes (outside the
|
|
// AssetDatabase). Refresh so Unity registers it before we import it,
|
|
// mirroring ImportModelFile. Skipped under the test import seam.
|
|
if (ImportOverrideForTests == null) AssetDatabase.Refresh();
|
|
AssetGenJob imported = r.ImportFn(r.Job, r.LocalPath);
|
|
if (imported != null) r.Job = imported;
|
|
if (r.Job.State != AssetGenJobState.Failed)
|
|
{
|
|
r.Job.State = AssetGenJobState.Done;
|
|
r.Job.Progress = 1f;
|
|
}
|
|
Persist(r.Job);
|
|
Finalize(r);
|
|
break;
|
|
}
|
|
}
|
|
catch (Exception e)
|
|
{
|
|
Fail(r, SecretRedactor.Scrub(e.Message));
|
|
}
|
|
}
|
|
|
|
private static string WriteFile(Runner r, byte[] bytes)
|
|
{
|
|
string chosen = !string.IsNullOrEmpty(r.OverrideExt) ? r.OverrideExt : r.Ext;
|
|
string ext = string.IsNullOrEmpty(chosen) ? "bin" : chosen.TrimStart('.').ToLowerInvariant();
|
|
string requestedRoot = !string.IsNullOrEmpty(r.OutputFolder) ? r.OutputFolder
|
|
: (AssetGenPrefs.OutputRoot + "/" + r.Subfolder);
|
|
if (!AssetGenPaths.TryGetAssetsFolder(requestedRoot, out string root))
|
|
root = AssetGenPrefs.DefaultOutputRoot + "/" + r.Subfolder;
|
|
string absRoot = AssetGenPaths.ToAbsolute(root);
|
|
Directory.CreateDirectory(absRoot);
|
|
string baseName = SanitizeName(r.Name);
|
|
string fileName = baseName + "." + ext;
|
|
string abs = Path.Combine(absRoot, fileName);
|
|
int n = 1;
|
|
while (File.Exists(abs)) { fileName = baseName + "_" + n++ + "." + ext; abs = Path.Combine(absRoot, fileName); }
|
|
File.WriteAllBytes(abs, bytes);
|
|
return (root.TrimEnd('/') + "/" + fileName).Replace('\\', '/');
|
|
}
|
|
|
|
private static string NameFrom(string explicitName, string prompt, string jobId)
|
|
{
|
|
if (!string.IsNullOrWhiteSpace(explicitName)) return explicitName;
|
|
if (!string.IsNullOrWhiteSpace(prompt)) return prompt;
|
|
return "asset_" + jobId.Substring(0, 8);
|
|
}
|
|
|
|
private static string SanitizeName(string raw)
|
|
{
|
|
if (string.IsNullOrWhiteSpace(raw)) return "asset";
|
|
var sb = new System.Text.StringBuilder();
|
|
foreach (char c in raw.Trim())
|
|
{
|
|
if (char.IsLetterOrDigit(c) || c == '_' || c == '-') sb.Append(c);
|
|
else if (c == ' ') sb.Append('_');
|
|
if (sb.Length >= 48) break;
|
|
}
|
|
string s = sb.ToString().Trim('_', '-');
|
|
return string.IsNullOrEmpty(s) ? "asset" : s;
|
|
}
|
|
|
|
private static void Fail(Runner r, string message)
|
|
{
|
|
r.Job.State = AssetGenJobState.Failed;
|
|
r.Job.Error = SecretRedactor.Scrub(string.IsNullOrEmpty(message) ? "Generation failed." : message);
|
|
Persist(r.Job);
|
|
Finalize(r);
|
|
}
|
|
|
|
private static void Finalize(Runner r)
|
|
{
|
|
Runners.Remove(r.Job.JobId);
|
|
try { r.Cts?.Dispose(); } catch { }
|
|
}
|
|
|
|
private static AssetGenJob NewJob(string kind, string provider, string action)
|
|
{
|
|
return new AssetGenJob
|
|
{
|
|
JobId = Guid.NewGuid().ToString("N"),
|
|
Kind = kind,
|
|
Provider = provider,
|
|
Action = action,
|
|
State = AssetGenJobState.Queued,
|
|
Progress = 0f,
|
|
};
|
|
}
|
|
|
|
private static void Persist(AssetGenJob job)
|
|
{
|
|
try
|
|
{
|
|
SessionState.SetString(JobKeyPrefix + job.JobId, JsonConvert.SerializeObject(job));
|
|
string index = SessionState.GetString(JobIndexKey, string.Empty);
|
|
var ids = new HashSet<string>(index.Split(new[] { ',' }, StringSplitOptions.RemoveEmptyEntries));
|
|
if (ids.Add(job.JobId))
|
|
SessionState.SetString(JobIndexKey, string.Join(",", ids));
|
|
}
|
|
catch { /* persistence is best-effort */ }
|
|
}
|
|
|
|
private static bool Faulted(Task t, out string error)
|
|
{
|
|
if (t.IsFaulted)
|
|
{
|
|
Exception ex = t.Exception?.GetBaseException();
|
|
error = SecretRedactor.Scrub(ex?.Message ?? "request failed");
|
|
return true;
|
|
}
|
|
if (t.IsCanceled) { error = "Canceled."; return true; }
|
|
error = null;
|
|
return false;
|
|
}
|
|
|
|
private static bool IsTerminal(AssetGenJobState s)
|
|
=> s == AssetGenJobState.Done || s == AssetGenJobState.Failed || s == AssetGenJobState.Canceled;
|
|
|
|
/// <summary>Only http(s) download URLs are allowed; provider responses are untrusted.</summary>
|
|
private static bool IsAllowedDownloadUrl(string url)
|
|
=> Uri.TryCreate(url, UriKind.Absolute, out Uri u)
|
|
&& (u.Scheme == Uri.UriSchemeHttps || u.Scheme == Uri.UriSchemeHttp);
|
|
|
|
private static double Now() => EditorApplication.timeSinceStartup;
|
|
|
|
internal static void ResetForTests()
|
|
{
|
|
foreach (var r in new List<Runner>(Runners.Values))
|
|
{
|
|
try { r.Cts?.Cancel(); r.Cts?.Dispose(); } catch { }
|
|
}
|
|
Runners.Clear();
|
|
string index = SessionState.GetString(JobIndexKey, string.Empty);
|
|
foreach (string id in index.Split(new[] { ',' }, StringSplitOptions.RemoveEmptyEntries))
|
|
SessionState.EraseString(JobKeyPrefix + id);
|
|
SessionState.EraseString(JobIndexKey);
|
|
Jobs.Clear();
|
|
TransportOverrideForTests = null;
|
|
ImportOverrideForTests = null;
|
|
PollIntervalSeconds = 3.0;
|
|
TimeoutSeconds = 600.0;
|
|
}
|
|
}
|
|
}
|