diff --git a/agent/aop_emit.go b/agent/aop_emit.go index 6164f5d6..019c512a 100644 --- a/agent/aop_emit.go +++ b/agent/aop_emit.go @@ -102,6 +102,7 @@ func (e *aopEmitter) sessionStart(model string) { event := &aop.Event{Payload: &aop.Event_SessionStarted{SessionStarted: &aop.SessionStarted{ Model: model, ParentSessionId: e.parentSessionID, ParentToolCallId: e.parentToolCallID, }}} + _ = types.SetSessionHistory(event, &types.SessionHistory{Mode: types.SessionHistory_MODE_INHERIT}) if e.delegation != nil { e.emitWithExt(event, e.delegation) return diff --git a/core/config/options.go b/core/config/options.go index 9787b91d..3f81e281 100644 --- a/core/config/options.go +++ b/core/config/options.go @@ -4,6 +4,7 @@ import ( "fmt" "io" "os" + "runtime" "strings" "github.com/chainreactors/aiscan/skills" @@ -241,6 +242,12 @@ func ResolvePrompt(value string) (string, error) { return prompt, nil } if err != nil { + // Natural-language prompts frequently contain punctuation that is not + // legal in a Windows filename (for example `host:port`). An invalid + // filename is evidence that this is text, not a prompt-file request. + if runtime.GOOS == "windows" && strings.ContainsAny(prompt, `<>:"|?*`) { + return prompt, nil + } return "", fmt.Errorf("stat prompt file %s: %w", prompt, err) } if !info.Mode().IsRegular() { diff --git a/core/config/options_test.go b/core/config/options_test.go index 4437a2f4..c316b806 100644 --- a/core/config/options_test.go +++ b/core/config/options_test.go @@ -33,6 +33,17 @@ func TestResolvePromptKeepsMissingPathAsNaturalLanguage(t *testing.T) { } } +func TestResolvePromptKeepsWindowsInvalidFilenameAsNaturalLanguage(t *testing.T) { + prompt := "check host:port and report the result" + got, err := ResolvePrompt(prompt) + if err != nil { + t.Fatalf("ResolvePrompt() error = %v", err) + } + if got != prompt { + t.Fatalf("ResolvePrompt() = %q, want %q", got, prompt) + } +} + func TestResolveTaskLoadsPromptFileAndAppendsInputs(t *testing.T) { path := filepath.Join(t.TempDir(), "task.md") if err := os.WriteFile(path, []byte("inspect the exposed services"), 0o600); err != nil { diff --git a/core/output/jsonl.go b/core/output/jsonl.go index ab492dea..7f472abd 100644 --- a/core/output/jsonl.go +++ b/core/output/jsonl.go @@ -2,6 +2,7 @@ package output import ( "bufio" + "bytes" "fmt" "io" "os" @@ -15,7 +16,8 @@ import ( ) // ScanJSONL decodes the canonical append-only AOP event stream one line at a -// time. Empty and non-JSON lines are ignored; malformed AOP event lines fail. +// time. Blank lines are allowed; every non-blank line must be a complete AOP +// event with a session and payload. func ScanJSONL(path string, visit func(*aop.Event) error) error { file, err := os.Open(path) if err != nil { @@ -25,16 +27,19 @@ func ScanJSONL(path string, visit func(*aop.Event) error) error { scanner := bufio.NewScanner(file) scanner.Buffer(make([]byte, 0, 256*1024), 64*1024*1024) for scanner.Scan() { - line := scanner.Bytes() - if len(line) == 0 || line[0] != '{' { + line := bytes.TrimSpace(scanner.Bytes()) + if len(line) == 0 { continue } + if line[0] != '{' { + return fmt.Errorf("AOP JSONL contains a non-event line") + } event := new(aop.Event) if err := protojson.Unmarshal(line, event); err != nil { return fmt.Errorf("decode AOP JSONL event: %w", err) } - if event.SessionId == "" || event.Payload == nil { - continue + if event.Id == "" || event.SessionId == "" || event.Payload == nil { + return fmt.Errorf("AOP JSONL event is missing id, session_id or payload") } if visit != nil { if err := visit(event); err != nil { @@ -57,30 +62,23 @@ func ReadJSONL(path string) ([]*aop.Event, error) { return events, err } -// DefaultRecorderMaxBytes caps one session's JSONL file. Legitimate sessions -// stay in the megabytes; the cap is a last-line defense so a runaway event -// source (a retry loop once emitted turn lifecycle pairs at microsecond -// cadence and wrote 182GB) cannot fill the disk. It is a capacity guard, not -// content policy: no event inspection or deduplication happens here. -const DefaultRecorderMaxBytes int64 = 2 << 30 // 2 GiB - // JSONLRecorder is the single append-only subscriber for persisted AOP events. type JSONLRecorder struct { - mu sync.Mutex - file *os.File - path string - unsub func() - err error - maxBytes int64 - size int64 - limited bool + mu sync.Mutex + file *os.File + path string + // seen is scoped to this recorder. Event IDs are unique within a session, + // not necessarily across independent sessions, so the key includes both. + seen map[string]struct{} + unsub func() + err error } func NewJSONLRecorder(bus *eventbus.Bus[*aop.Event], path string) (*JSONLRecorder, error) { if bus == nil { return nil, fmt.Errorf("AOP event bus is required") } - recorder := &JSONLRecorder{maxBytes: DefaultRecorderMaxBytes} + recorder := &JSONLRecorder{seen: make(map[string]struct{})} if err := recorder.Switch(path); err != nil { return nil, err } @@ -126,31 +124,40 @@ func (r *JSONLRecorder) Switch(path string) error { if err != nil { return err } - // The file is opened in append mode, so a resumed session starts with its - // existing size already counted against the cap. - var size int64 - if info, statErr := file.Stat(); statErr == nil { - size = info.Size() - } r.mu.Lock() + defer r.mu.Unlock() + seen, err := loadJSONLIDs(clean) + if err != nil { + _ = file.Close() + return err + } old := r.file r.file = file r.path = clean - r.size = size - r.limited = r.maxBytes > 0 && size >= r.maxBytes - r.mu.Unlock() + r.seen = seen if old != nil { if err := old.Close(); err != nil { - r.mu.Lock() if r.err == nil { r.err = err } - r.mu.Unlock() } } return nil } +func loadJSONLIDs(path string) (map[string]struct{}, error) { + seen := make(map[string]struct{}) + if err := ScanJSONL(path, func(event *aop.Event) error { + if event.Id != "" { + seen[event.SessionId+"\x00"+event.Id] = struct{}{} + } + return nil + }); err != nil { + return nil, err + } + return seen, nil +} + func (r *JSONLRecorder) Path() string { if r == nil { return "" @@ -164,13 +171,8 @@ func (r *JSONLRecorder) Write(event *aop.Event) error { if r == nil || event == nil { return nil } - // Fast path once the cap has tripped: skip the marshal so an ongoing - // event storm costs almost nothing per dropped event. - r.mu.Lock() - limited := r.limited - r.mu.Unlock() - if limited { - return nil + if event.Id == "" || event.SessionId == "" || event.Payload == nil { + return fmt.Errorf("AOP JSONL event requires id, session_id and payload") } line, err := (protojson.MarshalOptions{UseProtoNames: true}).Marshal(event) if err != nil { @@ -182,28 +184,22 @@ func (r *JSONLRecorder) Write(event *aop.Event) error { if r.file == nil { return io.ErrClosedPipe } - if r.limited { + key := event.SessionId + "\x00" + event.Id + if _, exists := r.seen[key]; exists { return nil } - if r.maxBytes > 0 && r.size+int64(len(line)) > r.maxBytes { - r.limited = true - // Leave one non-JSON marker line as durable evidence of the - // truncation; ScanJSONL skips lines that do not start with '{', so - // the file stays readable and resumable. - marker := fmt.Sprintf("# aiscan: JSONL size limit reached (limit=%d bytes); subsequent events are dropped\n", r.maxBytes) - _, _ = r.file.Write([]byte(marker)) - // The bus subscription stores the first Write error, so this - // surfaces once through Close instead of once per dropped event. - return fmt.Errorf("AOP JSONL %s reached the %d-byte size limit; subsequent events are dropped", r.path, r.maxBytes) - } n, err := r.file.Write(line) if err == nil && n != len(line) { err = io.ErrShortWrite } - if err == nil { - r.size += int64(n) + if err != nil { + return err } - return err + if r.seen == nil { + r.seen = make(map[string]struct{}) + } + r.seen[key] = struct{}{} + return nil } func (r *JSONLRecorder) Close() error { diff --git a/core/output/jsonl_test.go b/core/output/jsonl_test.go index 42a47ef6..904a4cba 100644 --- a/core/output/jsonl_test.go +++ b/core/output/jsonl_test.go @@ -4,12 +4,12 @@ import ( "fmt" "os" "path/filepath" - "strings" "sync" "testing" aop "github.com/chainreactors/aiscan/aop" "github.com/chainreactors/aiscan/core/eventbus" + "google.golang.org/protobuf/proto" ) func TestJSONLRecorderWritesConcurrentEventsAsCompleteLines(t *testing.T) { @@ -77,74 +77,76 @@ func TestJSONLRecorderSwitchesFilesWithoutReplayingHistory(t *testing.T) { } } -func TestJSONLRecorderStopsWritingAtSizeLimit(t *testing.T) { - path := filepath.Join(t.TempDir(), "capped.jsonl") +func TestJSONLRecorderSkipsDuplicateEventIDWithinSession(t *testing.T) { + path := filepath.Join(t.TempDir(), "events.jsonl") bus := eventbus.New[*aop.Event]() recorder, err := NewJSONLRecorder(bus, path) if err != nil { t.Fatal(err) } - recorder.mu.Lock() - recorder.maxBytes = 512 - recorder.mu.Unlock() - - const emitted = 100 - for i := 0; i < emitted; i++ { - bus.Emit(jsonlTestMessage(fmt.Sprintf("event-%d", i))) + event := &aop.Event{Id: "event-retry", SessionId: "session-1", Payload: &aop.Event_Status{Status: &aop.Status{State: "ready"}}} + bus.Emit(event) + bus.Emit(proto.Clone(event).(*aop.Event)) + if err := recorder.Close(); err != nil { + t.Fatal(err) } + events, err := ReadJSONL(path) + if err != nil { + t.Fatal(err) + } + if len(events) != 1 { + t.Fatalf("JSONL events = %d, want 1", len(events)) + } +} - info, err := os.Stat(path) +func TestJSONLRecorderLoadsExistingEventIDsAcrossRestart(t *testing.T) { + path := filepath.Join(t.TempDir(), "events.jsonl") + event := jsonlTestMessage("persisted") + bus := eventbus.New[*aop.Event]() + first, err := NewJSONLRecorder(bus, path) if err != nil { t.Fatal(err) } - // The cap bounds the payload; only the one marker line may exceed it. - if info.Size() > 512+256 { - t.Fatalf("file size = %d, want bounded near 512", info.Size()) + bus.Emit(event) + if err := first.Close(); err != nil { + t.Fatal(err) } - data, err := os.ReadFile(path) + second, err := NewJSONLRecorder(bus, path) if err != nil { t.Fatal(err) } - if !strings.Contains(string(data), "size limit reached") { - t.Fatalf("truncation marker missing:\n%s", data) + bus.Emit(proto.Clone(event).(*aop.Event)) + if err := second.Close(); err != nil { + t.Fatal(err) } - // The marker is a non-JSON comment line: the file must stay readable. events, err := ReadJSONL(path) if err != nil { t.Fatal(err) } - if len(events) == 0 || len(events) >= emitted { - t.Fatalf("persisted events = %d, want partial prefix of %d", len(events), emitted) - } - if err := recorder.Close(); err == nil || !strings.Contains(err.Error(), "size limit") { - t.Fatalf("Close() error = %v, want size limit error", err) + if len(events) != 1 { + t.Fatalf("events after recorder restart = %d, want 1", len(events)) } } -func TestJSONLRecorderSwitchResetsSizeBudget(t *testing.T) { - dir := t.TempDir() - bus := eventbus.New[*aop.Event]() - recorder, err := NewJSONLRecorder(bus, filepath.Join(dir, "first.jsonl")) - if err != nil { +func TestScanJSONLRejectsNonEventLines(t *testing.T) { + path := filepath.Join(t.TempDir(), "invalid.jsonl") + if err := os.WriteFile(path, []byte("traffic\n"), 0o644); err != nil { t.Fatal(err) } - recorder.mu.Lock() - recorder.maxBytes = 256 - recorder.mu.Unlock() - for i := 0; i < 10; i++ { - bus.Emit(jsonlTestMessage(fmt.Sprintf("first-%d", i))) - } - second := filepath.Join(dir, "second.jsonl") - if err := recorder.Switch(second); err != nil { - t.Fatal(err) + if _, err := ReadJSONL(path); err == nil { + t.Fatal("ReadJSONL accepted a non-event line") } - bus.Emit(jsonlTestMessage("after-switch")) - events, err := ReadJSONL(second) +} + +func TestJSONLRecorderRejectsEventsWithoutIdentity(t *testing.T) { + path := filepath.Join(t.TempDir(), "events.jsonl") + recorder, err := NewJSONLRecorder(eventbus.New[*aop.Event](), path) if err != nil { t.Fatal(err) } - if len(events) != 1 || events[0].Id != "after-switch" { - t.Fatalf("second file events = %#v, want the post-switch event", events) + defer recorder.Close() + if err := recorder.Write(&aop.Event{SessionId: "session", Payload: &aop.Event_Status{Status: &aop.Status{State: "ready"}}}); err == nil { + t.Fatal("Write accepted an event without an id") } } diff --git a/core/output/render_test.go b/core/output/render_test.go index 889f57f5..bb5014e9 100644 --- a/core/output/render_test.go +++ b/core/output/render_test.go @@ -127,7 +127,11 @@ func TestEventRendererSeparatesContinuationSessionHeaders(t *testing.T) { } func renderEvent(event *aop.Event) *aop.Event { - event.Id = "e-1" + if message := event.GetMessage(); message != nil && message.Id != "" { + event.Id = "e-" + message.Id + } else { + event.Id = "e-session-started" + } event.SessionId = "session-1" event.TurnId = "turn-1" event.Emitter = "aiscan" diff --git a/core/resources/resources.go b/core/resources/resources.go index 1ca5a679..ed129dee 100644 --- a/core/resources/resources.go +++ b/core/resources/resources.go @@ -122,6 +122,14 @@ func Init(ctx context.Context, opts Options) (*Set, error) { set.FingersConfig = fingers.NewConfig() set.FingersConfig.FullFingers = finalFullFingers set.NeutronConfig = neutron.NewConfig().WithTemplates(finalTemplates) + // Neutron compiles each template's HTTP transport when NewEngine runs. + // Carry the caller's egress proxy into the config before compilation so + // embedded and remote templates cannot bypass the Runner Hub. The engine + // package still applies the process default for compatibility with callers + // that construct neutron commands directly. + if opts.Proxy != "" { + set.NeutronConfig.WithProxy(opts.Proxy) + } set.Fingers, err = fingers.NewEngineWithFingers(finalFullFingers) if err != nil { diff --git a/core/resources/resources_test.go b/core/resources/resources_test.go index e6b9b717..4de8ee7d 100644 --- a/core/resources/resources_test.go +++ b/core/resources/resources_test.go @@ -3,8 +3,11 @@ package resources import ( "bytes" "context" + "net" + nethttp "net/http" "strings" "testing" + "time" fingerresources "github.com/chainreactors/fingers/resources" gogopkg "github.com/chainreactors/gogo/v2/pkg" @@ -12,6 +15,63 @@ import ( zombiepkg "github.com/chainreactors/zombie/pkg" ) +func TestInitBindsNeutronProxyBeforeTemplateCompilation(t *testing.T) { + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer listener.Close() + hit := make(chan struct{}, 1) + go func() { + conn, acceptErr := listener.Accept() + if acceptErr == nil { + hit <- struct{}{} + _ = conn.Close() + } + }() + + set, err := Init(context.Background(), Options{Proxy: "http://" + listener.Addr().String()}) + if err != nil { + t.Fatalf("Init() error = %v", err) + } + if set.Fingers != nil { + t.Cleanup(func() { _ = set.Fingers.Close() }) + } + if set.Neutron != nil { + t.Cleanup(func() { _ = set.Neutron.Close() }) + } + if set.Neutron == nil || set.Neutron.Count() == 0 { + t.Fatal("neutron engine has no compiled templates") + } + for _, template := range set.Neutron.Get() { + if template == nil { + continue + } + for _, request := range template.GetRequests() { + if request == nil || request.GetHTTPClient() == nil { + continue + } + transport, ok := request.GetHTTPClient().Transport.(*nethttp.Transport) + if !ok || transport.DialContext == nil { + t.Fatalf("template %q transport = %#v, want proxy dialer", template.Id, request.GetHTTPClient().Transport) + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + conn, _ := transport.DialContext(ctx, "tcp", "example.invalid:80") + cancel() + if conn != nil { + _ = conn.Close() + } + select { + case <-hit: + case <-time.After(time.Second): + t.Fatal("compiled neutron template did not dial the configured proxy") + } + return + } + } + t.Fatal("compiled neutron templates contain no HTTP request") +} + func TestInitUsesAiscanEmbeddedResources(t *testing.T) { oldFingerPrePort := fingerresources.PrePort oldFingerPortData := cloneBytes(fingerresources.PortData) diff --git a/pkg/commands/bash.go b/pkg/commands/bash.go index 9b866c9e..c67aa268 100644 --- a/pkg/commands/bash.go +++ b/pkg/commands/bash.go @@ -675,27 +675,9 @@ func (t *BashTool) proxyEnv(ctx context.Context) []string { callID := coretool.InvocationFromContext(ctx).CallID proxy, ca = t.egressResolver(callID) } - if proxy == "" { - return nil - } - env := []string{ - "ALL_PROXY=" + proxy, "all_proxy=" + proxy, - "HTTP_PROXY=" + proxy, "http_proxy=" + proxy, - "HTTPS_PROXY=" + proxy, "https_proxy=" + proxy, - } - // Point common HTTP clients at the MITM hub CA so intercepted HTTPS is - // trusted. Tools that use the system pool or pin certs ignore these and - // degrade to CONNECT-metadata capture, which is acceptable. - if ca != "" { - env = append(env, - "CURL_CA_BUNDLE="+ca, - "SSL_CERT_FILE="+ca, - "NODE_EXTRA_CA_CERTS="+ca, - "REQUESTS_CA_BUNDLE="+ca, - "GIT_SSL_CAINFO="+ca, - ) - } - return env + // Point the same common proxy/CA surface at child processes that built-in + // tools consume through Execution.Env. + return EgressEnvironment(proxy, ca) } func (t *BashTool) startMonitor(info tmux.Info, targetInbox inbox.Inbox) { diff --git a/pkg/commands/egress.go b/pkg/commands/egress.go new file mode 100644 index 00000000..2694ea14 --- /dev/null +++ b/pkg/commands/egress.go @@ -0,0 +1,94 @@ +package commands + +import "strings" + +// Egress is the per-invocation outbound route injected by the Runner. +// ProxyURL carries the call-scoped Hub identity; CAPath is populated only when +// the Hub is actively intercepting HTTPS traffic. +type Egress struct { + ProxyURL string + CAPath string +} + +// ResolveExecutionEgress resolves the route for one command invocation. The +// execution environment is intentionally authoritative for call-scoped Runner +// state; fallbackProxy is only the command's startup default. +func ResolveExecutionEgress(execution *Execution, fallbackProxy string) Egress { + if execution == nil { + return ResolveEgress(nil, fallbackProxy) + } + return ResolveEgress(execution.Env, fallbackProxy) +} + +var ( + // proxyEnvNames preserves the conventional environment surface exposed to + // child processes. Keep both cases because Windows and POSIX callers differ + // in how they spell environment keys. + proxyEnvNames = []string{ + "ALL_PROXY", "all_proxy", + "HTTP_PROXY", "http_proxy", + "HTTPS_PROXY", "https_proxy", + } + proxyLookupOrder = []string{ + "ALL_PROXY", "all_proxy", + "HTTPS_PROXY", "https_proxy", + "HTTP_PROXY", "http_proxy", + } + caEnvNames = []string{ + "CURL_CA_BUNDLE", "SSL_CERT_FILE", "NODE_EXTRA_CA_CERTS", + "REQUESTS_CA_BUNDLE", "GIT_SSL_CAINFO", + } +) + +// ResolveEgress resolves one invocation's egress from an environment slice. +// The explicit order is independent of how the caller sorted or assembled its +// environment, and falls back to the command's startup proxy when no call +// scoped value is present. +func ResolveEgress(env []string, fallbackProxy string) Egress { + values := make(map[string]string, len(env)) + for _, item := range env { + key, value, ok := strings.Cut(item, "=") + if !ok { + continue + } + values[key] = value + } + + resolved := Egress{} + for _, key := range proxyLookupOrder { + if value := strings.TrimSpace(values[key]); value != "" { + resolved.ProxyURL = value + break + } + } + if resolved.ProxyURL == "" { + resolved.ProxyURL = strings.TrimSpace(fallbackProxy) + } + for _, key := range caEnvNames { + if value := strings.TrimSpace(values[key]); value != "" { + resolved.CAPath = value + break + } + } + return resolved +} + +// EgressEnvironment returns the environment entries used by both built-in +// tools and child shell commands. An empty proxy intentionally produces no +// proxy variables, preserving direct-command behavior for non-Runner callers. +func EgressEnvironment(proxyURL, caPath string) []string { + proxyURL = strings.TrimSpace(proxyURL) + if proxyURL == "" { + return nil + } + env := make([]string, 0, len(proxyEnvNames)+len(caEnvNames)) + for _, key := range proxyEnvNames { + env = append(env, key+"="+proxyURL) + } + if caPath = strings.TrimSpace(caPath); caPath != "" { + for _, key := range caEnvNames { + env = append(env, key+"="+caPath) + } + } + return env +} diff --git a/pkg/commands/egress_test.go b/pkg/commands/egress_test.go new file mode 100644 index 00000000..c29ed180 --- /dev/null +++ b/pkg/commands/egress_test.go @@ -0,0 +1,82 @@ +package commands + +import "testing" + +func TestResolveEgressUsesStableProxyPrecedence(t *testing.T) { + tests := []struct { + name string + env []string + fallback string + wantProxy string + wantCA string + }{ + { + name: "all proxy wins independent of env order", + env: []string{"HTTP_PROXY=http://http", "ALL_PROXY=http://all"}, + fallback: "http://startup", + wantProxy: "http://all", + }, + { + name: "https beats http", + env: []string{"HTTP_PROXY=http://http", "HTTPS_PROXY=http://https"}, + wantProxy: "http://https", + }, + { + name: "fallback", + env: []string{"ALL_PROXY=", "HTTPS_PROXY= "}, + fallback: " http://startup ", + wantProxy: "http://startup", + }, + { + name: "ca precedence", + env: []string{"GIT_SSL_CAINFO=/git.pem", "SSL_CERT_FILE=/ssl.pem", "CURL_CA_BUNDLE=/curl.pem"}, + wantCA: "/curl.pem", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := ResolveEgress(tt.env, tt.fallback) + if got.ProxyURL != tt.wantProxy || got.CAPath != tt.wantCA { + t.Fatalf("ResolveEgress() = %#v, want proxy=%q ca=%q", got, tt.wantProxy, tt.wantCA) + } + }) + } +} + +func TestResolveExecutionEgressUsesInvocationEnvironment(t *testing.T) { + execution := &Execution{ + Env: []string{"ALL_PROXY=http://call-scoped", "SSL_CERT_FILE=/call-ca.pem"}, + } + got := ResolveExecutionEgress(execution, "http://startup") + if got.ProxyURL != "http://call-scoped" || got.CAPath != "/call-ca.pem" { + t.Fatalf("ResolveExecutionEgress() = %#v, want call-scoped proxy and CA", got) + } + + got = ResolveExecutionEgress(&Execution{}, "http://startup") + if got.ProxyURL != "http://startup" { + t.Fatalf("ResolveExecutionEgress(empty env) = %#v, want startup fallback", got) + } +} + +func TestEgressEnvironmentUsesSharedSurface(t *testing.T) { + got := EgressEnvironment("http://hub", "/tmp/mitm-ca.pem") + want := []string{ + "ALL_PROXY=http://hub", "all_proxy=http://hub", + "HTTP_PROXY=http://hub", "http_proxy=http://hub", + "HTTPS_PROXY=http://hub", "https_proxy=http://hub", + "CURL_CA_BUNDLE=/tmp/mitm-ca.pem", "SSL_CERT_FILE=/tmp/mitm-ca.pem", + "NODE_EXTRA_CA_CERTS=/tmp/mitm-ca.pem", "REQUESTS_CA_BUNDLE=/tmp/mitm-ca.pem", + "GIT_SSL_CAINFO=/tmp/mitm-ca.pem", + } + if len(got) != len(want) { + t.Fatalf("EgressEnvironment() = %#v, want %#v", got, want) + } + for i := range want { + if got[i] != want[i] { + t.Fatalf("EgressEnvironment()[%d] = %q, want %q", i, got[i], want[i]) + } + } + if got := EgressEnvironment("", "/tmp/mitm-ca.pem"); got != nil { + t.Fatalf("empty proxy environment = %#v, want nil", got) + } +} diff --git a/pkg/node/agent.go b/pkg/node/agent.go index db60ced1..d311476b 100644 --- a/pkg/node/agent.go +++ b/pkg/node/agent.go @@ -64,19 +64,18 @@ func runRemoteAgent(ctx context.Context, option *cfg.Option, logger telemetry.Lo logger.Debugf("websocket transport connection to %s", dialURL) connection := connectionConfig{ - ServerURL: option.ServerURL, - Name: runner.ResolveIOANodeName(option), - Registry: application.Commands, - AgentSubscribe: rt.Subscribe, - Progress: application.Progress, - Logger: logger, - Chat: chatHandler, - AgentRuntime: rt, - NodeID: nodeID, - Runtime: runner.DefaultRuntimeInfo(), - Status: func() *aop.AgentStatus { return runner.AgentStatus(option, application) }, - Menu: func() []*types.CommandSpec { return runner.CommandCatalog(application) }, - PTYRouter: func() (*terminal.Router, error) { return NewPTYRouter(application.Commands), nil }, + ServerURL: option.ServerURL, + Name: runner.ResolveIOANodeName(option), + Registry: application.Commands, + Agent: rt, + Progress: application.Progress, + Logger: logger, + Chat: chatHandler, + NodeID: nodeID, + Runtime: runner.DefaultRuntimeInfo(), + Status: func() *aop.AgentStatus { return runner.AgentStatus(option, application) }, + Menu: func() []*types.CommandSpec { return runner.CommandCatalog(application) }, + PTYRouter: func() (*terminal.Router, error) { return NewPTYRouter(application.Commands), nil }, } _ = connect(ctx, connection) }() diff --git a/pkg/node/connection.go b/pkg/node/connection.go index c17015b2..3e01d8ce 100644 --- a/pkg/node/connection.go +++ b/pkg/node/connection.go @@ -8,13 +8,52 @@ import ( "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/runner" "github.com/chainreactors/aiscan/pkg/terminal" types "github.com/chainreactors/aiscan/pkg/types" ) const DefaultWSPath = "/api/aop/node/ws" +// agentEndpoint is the sole event ingress/egress point for a node connection. +// Agent runtimes implement the optional control method as well; tool-only +// nodes use the eventBusEndpoint adapter below. Keeping publication and +// subscription on one object prevents a terminal event from being sent both +// through the runtime bus and as a direct protocol reply. +type agentEndpoint interface { + Subscribe(func(*aop.Event)) func() + EmitEvent(*aop.Event) +} + +type agentControlEndpoint interface { + agentEndpoint + HandleEnvelope(context.Context, *aop.Envelope, func(*aop.Envelope)) bool +} + +type eventBusEndpoint struct { + bus *eventbus.Bus[*aop.Event] +} + +func newEventBusEndpoint(bus *eventbus.Bus[*aop.Event]) agentEndpoint { + if bus == nil { + bus = eventbus.New[*aop.Event]() + } + return &eventBusEndpoint{bus: bus} +} + +func (e *eventBusEndpoint) Subscribe(fn func(*aop.Event)) func() { + if e == nil || e.bus == nil { + return func() {} + } + return e.bus.Subscribe(fn) +} + +func (e *eventBusEndpoint) EmitEvent(event *aop.Event) { + if e == nil || e.bus == nil || event == nil { + return + } + e.bus.Emit(event) +} + type connectionConfig struct { ServerURL string WSPath string @@ -24,15 +63,14 @@ type connectionConfig struct { // JSONFrames switches the wire codec from binary protobuf to standard // ProtoJSON text frames (used by hubs that speak JSON, e.g. Cairn). - JSONFrames bool - Registry *commands.CommandRegistry - AgentSubscribe func(func(*aop.Event)) func() - Progress *eventbus.Bus[*toolpb.Progress] - Logger telemetry.Logger - Chat *chatAgentHandler - // AgentRuntime handles the AOP core/command namespaces directly via - // HandleEnvelope; nil on tool-only nodes, which reject chat messages. - AgentRuntime *runner.AgentRuntime + JSONFrames bool + Registry *commands.CommandRegistry + // Agent is the single owner of connection-side events. Implementations that + // also satisfy agentControlEndpoint handle core/command namespaces. + Agent agentEndpoint + Progress *eventbus.Bus[*toolpb.Progress] + Logger telemetry.Logger + Chat *chatAgentHandler NodeID string Runtime *aop.AgentRuntimeInfo Status func() *aop.AgentStatus diff --git a/pkg/node/proto_connection.go b/pkg/node/proto_connection.go index 50eeb809..9f43246f 100644 --- a/pkg/node/proto_connection.go +++ b/pkg/node/proto_connection.go @@ -194,6 +194,12 @@ func serveAgentConnection(ctx context.Context, cc connectionConfig, logger telem if cc.Registry == nil { return fmt.Errorf("command registry is nil") } + // Every connection gets one event endpoint. Tool-only callers may provide + // their own event bus; a missing endpoint is backed by a private bus so + // successful tool results never need a second direct-delivery path. + if cc.Agent == nil { + cc.Agent = newEventBusEndpoint(nil) + } hello, err := BuildHello(cc.Name, cc.Registry, cc.NodeID, cc.Runtime) if err != nil { return err @@ -291,8 +297,8 @@ func serveAgentConnection(ctx context.Context, cc connectionConfig, logger telem sealed := make(map[string]time.Time) stats := NewAgentStatsTracker() - if cc.AgentSubscribe != nil { - unsubscribe := cc.AgentSubscribe(func(event *aop.Event) { + if cc.Agent != nil { + unsubscribe := cc.Agent.Subscribe(func(event *aop.Event) { if next, changed := stats.Observe(event); changed { send("", &aop.ProtocolMessage{Message: &aop.ProtocolMessage_AgentStats{AgentStats: next}}) } @@ -446,8 +452,8 @@ func newAgentConnectionNamespaceMux( return nil, err } if err := mux.Register(&types.CommandProtocolMessage{}, func(ctx context.Context, envelope *aop.Envelope, _ protobuf.Message, _ aop.SendFunc) error { - if cc.AgentRuntime != nil { - cc.AgentRuntime.HandleEnvelope(ctx, envelope, sendEnvelope) + if agent, ok := cc.Agent.(agentControlEndpoint); ok { + agent.HandleEnvelope(ctx, envelope, sendEnvelope) return nil } send(envelope.GetId(), protocolFailure("OPERATION_FAILED", "command handler is unavailable")) @@ -542,8 +548,8 @@ func handleAgentCoreMessage( cancel() return } - if cc.AgentRuntime != nil { - cc.AgentRuntime.HandleEnvelope(ctx, envelope, sendEnvelope) + if agent, ok := cc.Agent.(agentControlEndpoint); ok { + agent.HandleEnvelope(ctx, envelope, sendEnvelope) return } send(envelope.GetId(), protocolFailure("OPERATION_FAILED", "chat handler is unavailable")) @@ -561,12 +567,9 @@ func handleAgentToolMessage(ctx context.Context, cc connectionConfig, envelope * if request.Call.Id == "" { request.Call.Id = operationID } - if cc.AgentRuntime != nil && request.Call.Id == operationID && strings.TrimSpace(request.Call.Name) != "" { - cc.AgentRuntime.EmitEvent(&aop.Event{ - SessionId: request.SessionId, TurnId: request.TurnId, Emitter: "aiscan.agent", - Payload: &aop.Event_ToolCall{ToolCall: protobuf.CloneOf(request.Call)}, - }) - } + // The hub is the canonical publisher for a remotely dispatched tool.call. + // This node only publishes the terminal result; synthesizing the call here + // would duplicate the hub's session timeline entry. taskCtx, taskCancel := context.WithCancel(ctx) trackOperation(operationsMu, operations, operationID, taskCancel) // seal closes this call's artifact window so the forwarding subscriber drops @@ -593,14 +596,14 @@ func handleAgentToolMessage(ctx context.Context, cc connectionConfig, envelope * fail(err.Error()) return } - // Seal ahead of EmitEvent too: the runtime publishes onto the same bus the - // forwarding subscriber reads, so on an agent node the terminal reaches the - // wire from inside EmitEvent, ahead of the send below. + // The endpoint is the single event source for the connection. Its + // subscriber forwards the terminal to the wire; do not send a second copy. seal() - if cc.AgentRuntime != nil { - cc.AgentRuntime.EmitEvent(event) + if cc.Agent == nil { + fail("agent event endpoint is unavailable") + return } - send(replyTo, &aop.ProtocolMessage{Message: &aop.ProtocolMessage_Event{Event: event}}) + cc.Agent.EmitEvent(event) }() } diff --git a/pkg/node/proto_connection_test.go b/pkg/node/proto_connection_test.go index 1a06ebd8..43bc5b55 100644 --- a/pkg/node/proto_connection_test.go +++ b/pkg/node/proto_connection_test.go @@ -21,13 +21,48 @@ import ( execpb "github.com/chainreactors/aiscan/aop/exec" filepb "github.com/chainreactors/aiscan/aop/file" toolpb "github.com/chainreactors/aiscan/aop/tool" + cfg "github.com/chainreactors/aiscan/core/config" + "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/telemetry" + coretool "github.com/chainreactors/aiscan/core/tool" "github.com/chainreactors/aiscan/pkg/commands" + "github.com/chainreactors/aiscan/pkg/runner" types "github.com/chainreactors/aiscan/pkg/types" "github.com/gorilla/websocket" protobuf "google.golang.org/protobuf/proto" ) +type singleDeliveryProbeTool struct{} + +func (singleDeliveryProbeTool) Name() string { return "single_delivery_probe" } + +func (singleDeliveryProbeTool) Description() string { return "test tool" } + +func (singleDeliveryProbeTool) Definition() *aop.ToolDefinition { + return coretool.Def("single_delivery_probe", "test tool", struct{}{}) +} + +func (singleDeliveryProbeTool) Execute(context.Context, string) (*coretool.Result, error) { + return coretool.TextResult("probe result"), nil +} + +type trackingAgentEndpoint struct { + bus *eventbus.Bus[*aop.Event] + subscribed *bool +} + +func (e *trackingAgentEndpoint) Subscribe(fn func(*aop.Event)) func() { + *e.subscribed = true + return e.bus.Subscribe(fn) +} + +func (e *trackingAgentEndpoint) EmitEvent(event *aop.Event) { e.bus.Emit(event) } + +type panicAgentEndpoint struct{} + +func (panicAgentEndpoint) Subscribe(func(*aop.Event)) func() { return func() {} } +func (panicAgentEndpoint) EmitEvent(*aop.Event) { panic("send event boom") } + type handshakeThenEOFStream struct { helloID string recvs int @@ -56,10 +91,7 @@ func TestServeAgentConnectionSubscribesBeforePublishingMenu(t *testing.T) { Name: "runner-1", NodeID: "runner-1", Registry: commands.NewRegistry(), - AgentSubscribe: func(func(*aop.Event)) func() { - subscribed = true - return func() {} - }, + Agent: &trackingAgentEndpoint{bus: eventbus.New[*aop.Event](), subscribed: &subscribed}, Menu: func() []*types.CommandSpec { menuCalled = true if !subscribed { @@ -95,7 +127,7 @@ func TestToolOperationPanicIsReportedAndCleanedUp(t *testing.T) { request := &toolpb.Call{Call: &aop.ToolCall{Id: "op-panic", Name: "missing", Arguments: arguments}} handleAgentToolMessage( context.Background(), - connectionConfig{Registry: commands.NewRegistry(), Logger: logger}, + connectionConfig{Registry: commands.NewRegistry(), Logger: logger, Agent: panicAgentEndpoint{}}, &aop.Envelope{Id: "op-panic"}, &toolpb.ProtocolMessage{Message: &toolpb.ProtocolMessage_Call{Call: request}}, send, &operationsMu, operations, make(map[string]time.Time), @@ -151,6 +183,125 @@ func TestCancelOperationSealsTheCallArtifactWindow(t *testing.T) { } } +func TestAgentRuntimeToolResultUsesSingleDeliveryPath(t *testing.T) { + ctx := context.Background() + app, err := runner.NewApp(ctx, runner.ApplicationConfig{ + SkipEngines: true, + Logger: telemetry.NopLogger(), + }) + if err != nil { + t.Fatal(err) + } + defer app.Close() + rt, err := runner.NewAgentRuntime(ctx, &cfg.Option{}, telemetry.NopLogger(), &runner.RuntimeConfig{ + ExistingApp: app, + ProviderOptional: true, + }) + if err != nil { + t.Fatal(err) + } + defer rt.Close() + + registry := commands.NewRegistry() + registry.RegisterTool(singleDeliveryProbeTool{}) + runtimeEvents := make(chan *aop.Event, 1) + var runtimeToolCalls atomic.Int32 + unsubscribe := rt.Subscribe(func(event *aop.Event) { + if event == nil { + return + } + if event.GetToolCall() != nil { + runtimeToolCalls.Add(1) + } + if event.GetToolResult() != nil { + runtimeEvents <- event + } + }) + defer unsubscribe() + directMessages := make(chan protobuf.Message, 2) + send := func(_ string, message protobuf.Message) { directMessages <- message } + arguments, err := aop.JSONValue(map[string]any{}) + if err != nil { + t.Fatal(err) + } + request := &toolpb.Call{Call: &aop.ToolCall{ + Id: "single-delivery-op", + Name: "single_delivery_probe", + Arguments: arguments, + }} + handleAgentToolMessage( + ctx, + connectionConfig{ + Registry: registry, + Logger: telemetry.NopLogger(), + Agent: rt, + }, + &aop.Envelope{Id: "single-delivery-op"}, + &toolpb.ProtocolMessage{Message: &toolpb.ProtocolMessage_Call{Call: request}}, + send, + &sync.Mutex{}, + make(map[string]context.CancelFunc), + make(map[string]time.Time), + ) + + select { + case event := <-runtimeEvents: + if got := event.GetToolResult().GetName(); got != "single_delivery_probe" { + t.Fatalf("runtime tool result name = %q", got) + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for runtime tool result") + } + select { + case message := <-directMessages: + t.Fatalf("tool result was sent directly in addition to runtime event: %T", message) + case <-time.After(100 * time.Millisecond): + } + if got := runtimeToolCalls.Load(); got != 0 { + t.Fatalf("remote tool request unexpectedly emitted %d tool.call events; the hub is the canonical source", got) + } +} + +func TestToolOnlyNodeToolResultUsesEndpointDelivery(t *testing.T) { + registry := commands.NewRegistry() + registry.RegisterTool(singleDeliveryProbeTool{}) + wireEvents := make(chan *aop.Event, 1) + directMessages := make(chan protobuf.Message, 1) + endpoint := newEventBusEndpoint(nil) + endpoint.Subscribe(func(event *aop.Event) { + wireEvents <- event + }) + arguments, err := aop.JSONValue(map[string]any{}) + if err != nil { + t.Fatal(err) + } + handleAgentToolMessage( + context.Background(), + connectionConfig{Registry: registry, Logger: telemetry.NopLogger(), Agent: endpoint}, + &aop.Envelope{Id: "tool-only-op"}, + &toolpb.ProtocolMessage{Message: &toolpb.ProtocolMessage_Call{Call: &toolpb.Call{Call: &aop.ToolCall{ + Id: "tool-only-op", Name: "single_delivery_probe", Arguments: arguments, + }}}}, + func(_ string, message protobuf.Message) { directMessages <- message }, + &sync.Mutex{}, + make(map[string]context.CancelFunc), + make(map[string]time.Time), + ) + select { + case event := <-wireEvents: + if event.GetToolResult() == nil || event.GetToolResult().GetName() != "single_delivery_probe" { + t.Fatalf("endpoint tool result = %+v", event) + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for endpoint tool result") + } + select { + case message := <-directMessages: + t.Fatalf("tool result bypassed the endpoint: %T", message) + case <-time.After(100 * time.Millisecond): + } +} + func TestExecRequestCompletesWithOutput(t *testing.T) { command := "printf hello" if runtime.GOOS == "windows" { diff --git a/pkg/node/toolnode.go b/pkg/node/toolnode.go index 22000766..4e4515a4 100644 --- a/pkg/node/toolnode.go +++ b/pkg/node/toolnode.go @@ -84,17 +84,13 @@ func RunToolNode(ctx context.Context, cfg ToolNodeConfig) error { skillStore, _ := skills.LoadEmbeddedStore() menu = func() []*types.CommandSpec { return runner.RegistryCommandCatalog(cfg.Registry, skillStore) } } - var subscribe func(func(*aop.Event)) func() - if cfg.Events != nil { - subscribe = cfg.Events.Subscribe - } return connect(ctx, connectionConfig{ ServerURL: cfg.ServerURL, WSPath: cfg.WSPath, Name: runnerID, Token: cfg.Token, Registry: cfg.Registry, - AgentSubscribe: subscribe, + Agent: newEventBusEndpoint(cfg.Events), Progress: cfg.Progress, Logger: logger, NodeID: runnerID, @@ -109,7 +105,8 @@ func RunToolNode(ctx context.Context, cfg ToolNodeConfig) error { } // attachToolProgress forwards ephemeral progress onto the tool protocol. Raw -// artifacts travel as canonical AOP extension events through AgentSubscribe. +// artifacts travel as canonical AOP extension events through the Agent +// endpoint's single event stream. func attachToolProgress(progressBus *eventbus.Bus[*toolpb.Progress], send func(string, protobuf.Message)) func() { if progressBus == nil { return nil diff --git a/pkg/runner/runner.go b/pkg/runner/runner.go index 152f5f06..f5000887 100644 --- a/pkg/runner/runner.go +++ b/pkg/runner/runner.go @@ -22,6 +22,7 @@ import ( coretool "github.com/chainreactors/aiscan/core/tool" cmdpkg "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/pkg/tui" + types "github.com/chainreactors/aiscan/pkg/types" "github.com/chainreactors/aiscan/skills" "github.com/chainreactors/aiscan/tools/toolargs" ioaclient "github.com/chainreactors/ioa/client" @@ -101,6 +102,26 @@ func samePath(left, right string) bool { return left == right } +func validateFreshJSONLOutput(option *cfg.Option) error { + if option == nil || strings.TrimSpace(option.Resume) != "" || strings.TrimSpace(option.OutputFile) == "" { + return nil + } + info, err := os.Stat(option.OutputFile) + if os.IsNotExist(err) { + return nil + } + if err != nil { + return fmt.Errorf("stat AOP JSONL output %s: %w", option.OutputFile, err) + } + if info.IsDir() { + return fmt.Errorf("AOP JSONL output %s is a directory", option.OutputFile) + } + if info.Size() > 0 { + return fmt.Errorf("AOP JSONL output %s already exists and is not empty; use --resume to append an existing session", option.OutputFile) + } + return nil +} + // RunOutput is the presentation sink an entry point may attach to a runtime. // The runtime never constructs one — CLI/TUI hosts inject it; headless hosts // (stdio, WebSocket nodes, the web hub) leave it nil. @@ -132,6 +153,10 @@ func NewAgentRuntime(ctx context.Context, option *cfg.Option, logger telemetry.L runtimeCancel() return nil, fmt.Errorf("agent runtime option is required") } + if err := validateFreshJSONLOutput(option); err != nil { + runtimeCancel() + return nil, err + } rt := &AgentRuntime{ ctx: runtimeCtx, cancel: runtimeCancel, @@ -157,7 +182,6 @@ func NewAgentRuntime(ctx context.Context, option *cfg.Option, logger telemetry.L rt.option.OutputFile = recordPath rt.recordPath = recordPath } - if rc != nil && rc.ExistingApp != nil { rt.app = rc.ExistingApp } else { @@ -650,7 +674,7 @@ func RunDirectScannerMode(ctx context.Context, option *cfg.Option, rest []string } startedAt := time.Now() if application.Events != nil { - application.Events.sessionStarted(sessionID, emitter, &aop.SessionStarted{}) + application.Events.sessionStarted(sessionID, emitter, &aop.SessionStarted{}, types.SessionHistory_MODE_INHERIT) application.Events.Emit(&aop.Event{ SessionId: sessionID, TurnId: turnID, Emitter: emitter, Payload: &aop.Event_TurnStarted{TurnStarted: &aop.TurnStarted{}}, diff --git a/pkg/runner/runner_test.go b/pkg/runner/runner_test.go index 86739035..b0910819 100644 --- a/pkg/runner/runner_test.go +++ b/pkg/runner/runner_test.go @@ -2,6 +2,7 @@ package runner import ( "context" + "os" "path/filepath" "strings" "testing" @@ -14,6 +15,7 @@ import ( "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/telemetry" + types "github.com/chainreactors/aiscan/pkg/types" "google.golang.org/protobuf/types/known/timestamppb" ) @@ -118,6 +120,62 @@ func TestResumeRestoresAndAppendsAOPStream(t *testing.T) { }) } +func TestContinuationReferencesHistoryWithoutReemittingLargeMessages(t *testing.T) { + path := filepath.Join(t.TempDir(), "continuation.jsonl") + option := &cfg.Option{MiscOptions: cfg.MiscOptions{OutputFile: path}} + provider := new(persistenceProvider) + app, runtime := newPersistenceRuntimeWithMode(t, option, provider, REPLEphemeral) + _ = app + + root, err := runtime.OpenSession(context.Background(), SessionOptions{ID: MainREPLName}) + if err != nil { + t.Fatal(err) + } + large := strings.Repeat("x", 4<<20) + oldID := root.ID() + runtime.sessionEvents.Emit(&aop.Event{ + SessionId: root.ID(), TurnId: "turn-1", Emitter: "aiscan", + Payload: &aop.Event_Message{Message: &aop.Message{Id: "m-1", Role: "user", Content: []*aop.Content{aop.Text(large)}}}, + }) + before, err := os.Stat(path) + if err != nil { + t.Fatal(err) + } + + if _, err := root.rotate(context.Background(), SessionCloseResumed, root.ID(), root.MessagesSnapshot(), ""); err != nil { + t.Fatal(err) + } + after, err := os.Stat(path) + if err != nil { + t.Fatal(err) + } + if growth := after.Size() - before.Size(); growth > 64<<10 { + t.Fatalf("continuation appended %d bytes for inherited history", growth) + } + + events, err := output.ReadJSONL(path) + if err != nil { + t.Fatal(err) + } + childID := root.ID() + if childID == oldID { + t.Fatalf("rotation did not create a child session: %q", childID) + } + for _, event := range events { + if event.SessionId == childID && (event.GetMessage() != nil || event.GetToolResult() != nil) { + t.Fatalf("inherited history was re-emitted in child stream: %s", aop.Kind(event)) + } + } + + data, err := loadResumeState(path) + if err != nil { + t.Fatal(err) + } + if len(data.Messages) != 1 || data.Messages[0].Content[0].GetText().GetText() != large { + t.Fatalf("resumed inherited history = %d messages, want the original large message", len(data.Messages)) + } +} + func TestREPLResumeLoadsMainSessionContext(t *testing.T) { resumePath := filepath.Join(t.TempDir(), "repl-resume.jsonl") writePersistenceSessionForID(t, resumePath, MainREPLName) @@ -251,7 +309,7 @@ func TestCompactRotatesAndPersistsOnlyCompactedContext(t *testing.T) { } } -func TestInteractiveResumeRotatesAndBootstrapsSelectedContext(t *testing.T) { +func TestInteractiveResumeRotatesAndUsesSelectedContext(t *testing.T) { dir := t.TempDir() currentPath := filepath.Join(dir, "current.jsonl") resumePath := filepath.Join(dir, "selected.jsonl") @@ -304,6 +362,19 @@ func TestInteractiveResumeRotatesAndBootstrapsSelectedContext(t *testing.T) { } } +func TestFreshJSONLOutputRejectsNonEmptyExistingFile(t *testing.T) { + path := filepath.Join(t.TempDir(), "existing.jsonl") + if err := os.WriteFile(path, []byte("{}\n"), 0o644); err != nil { + t.Fatal(err) + } + if err := validateFreshJSONLOutput(&cfg.Option{MiscOptions: cfg.MiscOptions{OutputFile: path}}); err == nil { + t.Fatal("non-empty output file was accepted without --resume") + } + if err := validateFreshJSONLOutput(&cfg.Option{AgentOptions: cfg.AgentOptions{Resume: path}}); err != nil { + t.Fatalf("resume output was rejected: %v", err) + } +} + func assertRotationEvents(t *testing.T, events []*aop.Event, oldID, newID, reason string) { t.Helper() var ended, started bool @@ -375,6 +446,7 @@ func writePersistenceSessionForID(t *testing.T, path, sessionID string) { {Id: "e-3", EmittedAt: timestamp, SessionId: sessionID, TurnId: "old-turn", Emitter: "aiscan", Seq: 3, Payload: &aop.Event_Message{Message: &aop.Message{Id: "m-2", Role: "assistant", Content: []*aop.Content{aop.Text("old assistant")}}}}, {Id: "e-4", EmittedAt: timestamp, SessionId: sessionID, Emitter: "aiscan", Seq: 4, Payload: &aop.Event_SessionEnded{SessionEnded: &aop.SessionEnded{Reason: "completed"}}}, } + _ = types.SetSessionHistory(events[0], &types.SessionHistory{Mode: types.SessionHistory_MODE_INHERIT}) bus := eventbus.New[*aop.Event]() writer, err := output.NewJSONLRecorder(bus, path) if err != nil { diff --git a/pkg/runner/runtime_session.go b/pkg/runner/runtime_session.go index 41e2e906..d4d1c3a9 100644 --- a/pkg/runner/runtime_session.go +++ b/pkg/runner/runtime_session.go @@ -37,6 +37,11 @@ type SessionOptions struct { ParentToolCallID string AgentName string Messages []*aop.Message + // HistorySnapshot marks Messages as a new persisted transcript snapshot. + // Ordinary continuations keep the in-memory context and refer to their + // parent session instead; replaying those messages as events would append + // every large tool result again to JSONL and the durable event stream. + HistorySnapshot bool } type SessionCloseReason string @@ -154,8 +159,12 @@ func (e *sessionEmitter) Emit(event *aop.Event) { e.bus.Emit(event) } -func (e *sessionEmitter) sessionStarted(sessionID, agentName string, started *aop.SessionStarted) { - e.Emit(&aop.Event{SessionId: sessionID, Emitter: agentName, Payload: &aop.Event_SessionStarted{SessionStarted: started}}) +func (e *sessionEmitter) sessionStarted(sessionID, agentName string, started *aop.SessionStarted, historyMode types.SessionHistory_Mode) { + event := &aop.Event{SessionId: sessionID, Emitter: agentName, Payload: &aop.Event_SessionStarted{SessionStarted: started}} + if historyMode != types.SessionHistory_MODE_UNSPECIFIED { + _ = types.SetSessionHistory(event, &types.SessionHistory{Mode: historyMode}) + } + e.Emit(event) } func (e *sessionEmitter) sessionEnded(sessionID, agentName, reason string) { @@ -646,10 +655,14 @@ func (rt *AgentRuntime) OpenSession(ctx context.Context, options SessionOptions) }) } go rt.runSession(state) + historyMode := types.SessionHistory_MODE_INHERIT + if options.HistorySnapshot { + historyMode = types.SessionHistory_MODE_SNAPSHOT + } rt.sessionEvents.sessionStarted(id, agentName, &aop.SessionStarted{ Model: rt.config.Model, ParentSessionId: options.ParentSessionID, ParentToolCallId: options.ParentToolCallID, - }) - if options.ParentSessionID != "" && options.ParentToolCallID == "" && len(options.Messages) > 0 { + }, historyMode) + if options.HistorySnapshot && len(options.Messages) > 0 { emitContinuationMessages(state, prepareContinuationMessages(options.Messages)) } return public, nil @@ -1036,7 +1049,7 @@ func (s *Session) rotate(ctx context.Context, reason SessionCloseReason, parentS newID := rt.nextContinuationID(logicalID) continuation, err := rt.OpenSession(ctx, SessionOptions{ ID: newID, LogicalID: logicalID, ParentSessionID: parentSessionID, - AgentName: agentName, Messages: prepared, + AgentName: agentName, Messages: prepared, HistorySnapshot: reason == SessionCloseCompacted, }) if err != nil { return nil, err diff --git a/pkg/runner/session_jsonl.go b/pkg/runner/session_jsonl.go index 7eb6e681..65bb58fb 100644 --- a/pkg/runner/session_jsonl.go +++ b/pkg/runner/session_jsonl.go @@ -32,12 +32,23 @@ type resumeStream struct { messageCounter int64 order int started bool + closedReason string + historyMode types.SessionHistory_Mode } func loadResumeState(path string) (*resumeState, error) { streams := make(map[string]*resumeStream) + seenEventIDs := make(map[string]struct{}) order := 0 err := output.ScanJSONL(path, func(event *aop.Event) error { + if event.Id == "" { + return fmt.Errorf("event in %s has no id", path) + } + eventKey := event.SessionId + "\x00" + event.Id + if _, exists := seenEventIDs[eventKey]; exists { + return fmt.Errorf("event id %s is duplicated in session %s", event.Id, event.SessionId) + } + seenEventIDs[eventKey] = struct{}{} stream := streams[event.SessionId] if stream == nil { order++ @@ -49,9 +60,19 @@ func loadResumeState(path string) (*resumeState, error) { stream.started = true stream.parentID = payload.SessionStarted.ParentSessionId stream.parentToolCall = payload.SessionStarted.ParentToolCallId + history, ok, err := types.GetSessionHistory(event) + if err != nil { + return fmt.Errorf("session %s has invalid history metadata: %w", event.SessionId, err) + } + if !ok || history.GetMode() == types.SessionHistory_MODE_UNSPECIFIED { + return fmt.Errorf("session %s has no explicit history metadata", event.SessionId) + } + stream.historyMode = history.GetMode() if payload.SessionStarted.Model != "" { stream.model = payload.SessionStarted.Model } + case *aop.Event_SessionEnded: + stream.closedReason = payload.SessionEnded.Reason case *aop.Event_Message: if payload.Message == nil || (payload.Message.Role != "user" && payload.Message.Role != "assistant") { return nil @@ -88,12 +109,67 @@ func loadResumeState(path string) (*resumeState, error) { if selected == nil { return nil, fmt.Errorf("no resumable AOP session found in %s", path) } + messages, counter, err := resumeStreamMessages(selected, streams) + if err != nil { + return nil, err + } return &resumeState{ - SessionID: selected.id, Model: selected.model, Messages: selected.messages, - MessageCounter: selected.messageCounter, + SessionID: selected.id, Model: selected.model, Messages: messages, + MessageCounter: counter, }, nil } +// resumeStreamMessages reconstructs the in-memory transcript without creating +// new events for inherited history. A compacted child explicitly declares a +// snapshot and supersedes its parent; all other sessions inherit their parent +// transcript and only contribute their own turn messages. +func resumeStreamMessages(selected *resumeStream, streams map[string]*resumeStream) ([]*aop.Message, int64, error) { + if selected == nil { + return nil, 0, nil + } + chain := make([]*resumeStream, 0, 4) + seen := make(map[string]struct{}) + current := selected + for current != nil { + if _, ok := seen[current.id]; ok { + return nil, 0, fmt.Errorf("session parent cycle detected at %s", current.id) + } + seen[current.id] = struct{}{} + chain = append(chain, current) + if current.historyMode == types.SessionHistory_MODE_SNAPSHOT || current.parentID == "" || current.parentToolCall != "" { + break + } + // /clear and /compact deliberately reset or replace the parent context; + // do not resurrect the discarded history when loading the file later. + if parent := streams[current.parentID]; parent != nil { + if parent.closedReason == string(SessionCloseCleared) || parent.closedReason == string(SessionCloseCompacted) { + break + } + if !parent.started { + return nil, 0, fmt.Errorf("session %s refers to parent %s without a session.started event", current.id, current.parentID) + } + } else { + return nil, 0, fmt.Errorf("session %s refers to missing parent %s", current.id, current.parentID) + } + current = streams[current.parentID] + } + + var messages []*aop.Message + var counter int64 + for i := len(chain) - 1; i >= 0; i-- { + stream := chain[i] + for _, message := range stream.messages { + if message == nil { + continue + } + messages = append(messages, proto.CloneOf(message)) + counter = max(counter, messageIDSequence(message.Id)) + } + counter = max(counter, stream.messageCounter) + } + return messages, counter, nil +} + func messageIDSequence(id string) int64 { if !strings.HasPrefix(id, "m-") { return 0 diff --git a/pkg/runner/session_jsonl_test.go b/pkg/runner/session_jsonl_test.go index 2df8e0f5..06620e75 100644 --- a/pkg/runner/session_jsonl_test.go +++ b/pkg/runner/session_jsonl_test.go @@ -10,6 +10,8 @@ import ( toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" + types "github.com/chainreactors/aiscan/pkg/types" + "google.golang.org/protobuf/encoding/protojson" "google.golang.org/protobuf/types/known/anypb" "google.golang.org/protobuf/types/known/timestamppb" ) @@ -60,8 +62,55 @@ func TestListSavedSessionsOnlyReadsJSONL(t *testing.T) { } } +func TestLoadResumeStateRejectsEventsWithoutHistoryMetadata(t *testing.T) { + path := filepath.Join(t.TempDir(), "unversioned.jsonl") + writeSessionEvents(t, path, []*aop.Event{ + {Id: "start", SessionId: "root", Payload: &aop.Event_SessionStarted{SessionStarted: &aop.SessionStarted{}}}, + sessionTestEvent("root", &aop.Event{Payload: &aop.Event_Message{Message: &aop.Message{Id: "m-1", Role: "user", Content: []*aop.Content{aop.Text("hello")}}}}), + }) + if _, err := loadResumeState(path); err == nil { + t.Fatal("loadResumeState accepted an event stream without explicit history metadata") + } +} + +func TestLoadResumeStateRejectsDuplicateEventIDs(t *testing.T) { + path := filepath.Join(t.TempDir(), "duplicate.jsonl") + start := sessionTestEvent("root", &aop.Event{Payload: &aop.Event_SessionStarted{SessionStarted: &aop.SessionStarted{}}}) + start.Id = "same-event" + message := sessionTestEvent("root", &aop.Event{Payload: &aop.Event_Message{Message: &aop.Message{Id: "m-1", Role: "user", Content: []*aop.Content{aop.Text("hello")}}}}) + message.Id = "same-event" + marshal := protojson.MarshalOptions{UseProtoNames: true} + first, err := marshal.Marshal(start) + if err != nil { + t.Fatal(err) + } + second, err := marshal.Marshal(message) + if err != nil { + t.Fatal(err) + } + data := append(append(first, '\n'), append(second, '\n')...) + if err := os.WriteFile(path, data, 0o644); err != nil { + t.Fatal(err) + } + if _, err := loadResumeState(path); err == nil { + t.Fatal("loadResumeState accepted duplicate event IDs") + } +} + func sessionTestEvent(sessionID string, event *aop.Event) *aop.Event { - event.Id = "event" + switch payload := event.Payload.(type) { + case *aop.Event_SessionStarted: + event.Id = "event-session-started-" + sessionID + _ = types.SetSessionHistory(event, &types.SessionHistory{Mode: types.SessionHistory_MODE_INHERIT}) + case *aop.Event_Message: + event.Id = "event-message-" + payload.Message.Id + case *aop.Event_ToolResult: + event.Id = "event-tool-result-" + payload.ToolResult.CallId + case *aop.Event_Extension: + event.Id = "event-extension-" + sessionID + default: + event.Id = "event-" + sessionID + } event.SessionId = sessionID event.TurnId = "turn-1" event.Emitter = "aiscan" diff --git a/pkg/types/chat.pb.go b/pkg/types/chat.pb.go index fcbae1e2..d6e77303 100644 --- a/pkg/types/chat.pb.go +++ b/pkg/types/chat.pb.go @@ -23,6 +23,101 @@ const ( _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) ) +type SessionHistory_Mode int32 + +const ( + SessionHistory_MODE_UNSPECIFIED SessionHistory_Mode = 0 + SessionHistory_MODE_INHERIT SessionHistory_Mode = 1 + SessionHistory_MODE_SNAPSHOT SessionHistory_Mode = 2 +) + +// Enum value maps for SessionHistory_Mode. +var ( + SessionHistory_Mode_name = map[int32]string{ + 0: "MODE_UNSPECIFIED", + 1: "MODE_INHERIT", + 2: "MODE_SNAPSHOT", + } + SessionHistory_Mode_value = map[string]int32{ + "MODE_UNSPECIFIED": 0, + "MODE_INHERIT": 1, + "MODE_SNAPSHOT": 2, + } +) + +func (x SessionHistory_Mode) Enum() *SessionHistory_Mode { + p := new(SessionHistory_Mode) + *p = x + return p +} + +func (x SessionHistory_Mode) String() string { + return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x)) +} + +func (SessionHistory_Mode) Descriptor() protoreflect.EnumDescriptor { + return file_types_chat_proto_enumTypes[0].Descriptor() +} + +func (SessionHistory_Mode) Type() protoreflect.EnumType { + return &file_types_chat_proto_enumTypes[0] +} + +func (x SessionHistory_Mode) Number() protoreflect.EnumNumber { + return protoreflect.EnumNumber(x) +} + +// Deprecated: Use SessionHistory_Mode.Descriptor instead. +func (SessionHistory_Mode) EnumDescriptor() ([]byte, []int) { + return file_types_chat_proto_rawDescGZIP(), []int{0, 0} +} + +// SessionHistory is persisted as an AOP event extension. It makes transcript +// inheritance explicit without changing the shared AOP protocol schema. +type SessionHistory struct { + state protoimpl.MessageState `protogen:"open.v1"` + Mode SessionHistory_Mode `protobuf:"varint,1,opt,name=mode,proto3,enum=aiscan.chat.SessionHistory_Mode" json:"mode,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *SessionHistory) Reset() { + *x = SessionHistory{} + mi := &file_types_chat_proto_msgTypes[0] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *SessionHistory) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*SessionHistory) ProtoMessage() {} + +func (x *SessionHistory) ProtoReflect() protoreflect.Message { + mi := &file_types_chat_proto_msgTypes[0] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use SessionHistory.ProtoReflect.Descriptor instead. +func (*SessionHistory) Descriptor() ([]byte, []int) { + return file_types_chat_proto_rawDescGZIP(), []int{0} +} + +func (x *SessionHistory) GetMode() SessionHistory_Mode { + if x != nil { + return x.Mode + } + return SessionHistory_MODE_UNSPECIFIED +} + type SessionRecord struct { state protoimpl.MessageState `protogen:"open.v1"` Session *aop.Session `protobuf:"bytes,1,opt,name=session,proto3" json:"session,omitempty"` @@ -36,7 +131,7 @@ type SessionRecord struct { func (x *SessionRecord) Reset() { *x = SessionRecord{} - mi := &file_types_chat_proto_msgTypes[0] + mi := &file_types_chat_proto_msgTypes[1] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -48,7 +143,7 @@ func (x *SessionRecord) String() string { func (*SessionRecord) ProtoMessage() {} func (x *SessionRecord) ProtoReflect() protoreflect.Message { - mi := &file_types_chat_proto_msgTypes[0] + mi := &file_types_chat_proto_msgTypes[1] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -61,7 +156,7 @@ func (x *SessionRecord) ProtoReflect() protoreflect.Message { // Deprecated: Use SessionRecord.ProtoReflect.Descriptor instead. func (*SessionRecord) Descriptor() ([]byte, []int) { - return file_types_chat_proto_rawDescGZIP(), []int{0} + return file_types_chat_proto_rawDescGZIP(), []int{1} } func (x *SessionRecord) GetSession() *aop.Session { @@ -110,7 +205,7 @@ type ListSessionsRequest struct { func (x *ListSessionsRequest) Reset() { *x = ListSessionsRequest{} - mi := &file_types_chat_proto_msgTypes[1] + mi := &file_types_chat_proto_msgTypes[2] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -122,7 +217,7 @@ func (x *ListSessionsRequest) String() string { func (*ListSessionsRequest) ProtoMessage() {} func (x *ListSessionsRequest) ProtoReflect() protoreflect.Message { - mi := &file_types_chat_proto_msgTypes[1] + mi := &file_types_chat_proto_msgTypes[2] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -135,7 +230,7 @@ func (x *ListSessionsRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use ListSessionsRequest.ProtoReflect.Descriptor instead. func (*ListSessionsRequest) Descriptor() ([]byte, []int) { - return file_types_chat_proto_rawDescGZIP(), []int{1} + return file_types_chat_proto_rawDescGZIP(), []int{2} } func (x *ListSessionsRequest) GetAfterCursor() string { @@ -169,7 +264,7 @@ type ListSessionsResponse struct { func (x *ListSessionsResponse) Reset() { *x = ListSessionsResponse{} - mi := &file_types_chat_proto_msgTypes[2] + mi := &file_types_chat_proto_msgTypes[3] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -181,7 +276,7 @@ func (x *ListSessionsResponse) String() string { func (*ListSessionsResponse) ProtoMessage() {} func (x *ListSessionsResponse) ProtoReflect() protoreflect.Message { - mi := &file_types_chat_proto_msgTypes[2] + mi := &file_types_chat_proto_msgTypes[3] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -194,7 +289,7 @@ func (x *ListSessionsResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use ListSessionsResponse.ProtoReflect.Descriptor instead. func (*ListSessionsResponse) Descriptor() ([]byte, []int) { - return file_types_chat_proto_rawDescGZIP(), []int{2} + return file_types_chat_proto_rawDescGZIP(), []int{3} } func (x *ListSessionsResponse) GetSessions() []*SessionRecord { @@ -220,7 +315,7 @@ type GetSessionRequest struct { func (x *GetSessionRequest) Reset() { *x = GetSessionRequest{} - mi := &file_types_chat_proto_msgTypes[3] + mi := &file_types_chat_proto_msgTypes[4] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -232,7 +327,7 @@ func (x *GetSessionRequest) String() string { func (*GetSessionRequest) ProtoMessage() {} func (x *GetSessionRequest) ProtoReflect() protoreflect.Message { - mi := &file_types_chat_proto_msgTypes[3] + mi := &file_types_chat_proto_msgTypes[4] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -245,7 +340,7 @@ func (x *GetSessionRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use GetSessionRequest.ProtoReflect.Descriptor instead. func (*GetSessionRequest) Descriptor() ([]byte, []int) { - return file_types_chat_proto_rawDescGZIP(), []int{3} + return file_types_chat_proto_rawDescGZIP(), []int{4} } func (x *GetSessionRequest) GetSessionId() string { @@ -264,7 +359,7 @@ type GetSessionResponse struct { func (x *GetSessionResponse) Reset() { *x = GetSessionResponse{} - mi := &file_types_chat_proto_msgTypes[4] + mi := &file_types_chat_proto_msgTypes[5] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -276,7 +371,7 @@ func (x *GetSessionResponse) String() string { func (*GetSessionResponse) ProtoMessage() {} func (x *GetSessionResponse) ProtoReflect() protoreflect.Message { - mi := &file_types_chat_proto_msgTypes[4] + mi := &file_types_chat_proto_msgTypes[5] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -289,7 +384,7 @@ func (x *GetSessionResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use GetSessionResponse.ProtoReflect.Descriptor instead. func (*GetSessionResponse) Descriptor() ([]byte, []int) { - return file_types_chat_proto_rawDescGZIP(), []int{4} + return file_types_chat_proto_rawDescGZIP(), []int{5} } func (x *GetSessionResponse) GetSession() *SessionRecord { @@ -311,7 +406,7 @@ type ResetSessionRequest struct { func (x *ResetSessionRequest) Reset() { *x = ResetSessionRequest{} - mi := &file_types_chat_proto_msgTypes[5] + mi := &file_types_chat_proto_msgTypes[6] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -323,7 +418,7 @@ func (x *ResetSessionRequest) String() string { func (*ResetSessionRequest) ProtoMessage() {} func (x *ResetSessionRequest) ProtoReflect() protoreflect.Message { - mi := &file_types_chat_proto_msgTypes[5] + mi := &file_types_chat_proto_msgTypes[6] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -336,7 +431,7 @@ func (x *ResetSessionRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use ResetSessionRequest.ProtoReflect.Descriptor instead. func (*ResetSessionRequest) Descriptor() ([]byte, []int) { - return file_types_chat_proto_rawDescGZIP(), []int{5} + return file_types_chat_proto_rawDescGZIP(), []int{6} } func (x *ResetSessionRequest) GetRequestId() string { @@ -377,7 +472,7 @@ type ResetSessionReceipt struct { func (x *ResetSessionReceipt) Reset() { *x = ResetSessionReceipt{} - mi := &file_types_chat_proto_msgTypes[6] + mi := &file_types_chat_proto_msgTypes[7] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -389,7 +484,7 @@ func (x *ResetSessionReceipt) String() string { func (*ResetSessionReceipt) ProtoMessage() {} func (x *ResetSessionReceipt) ProtoReflect() protoreflect.Message { - mi := &file_types_chat_proto_msgTypes[6] + mi := &file_types_chat_proto_msgTypes[7] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -402,7 +497,7 @@ func (x *ResetSessionReceipt) ProtoReflect() protoreflect.Message { // Deprecated: Use ResetSessionReceipt.ProtoReflect.Descriptor instead. func (*ResetSessionReceipt) Descriptor() ([]byte, []int) { - return file_types_chat_proto_rawDescGZIP(), []int{6} + return file_types_chat_proto_rawDescGZIP(), []int{7} } func (x *ResetSessionReceipt) GetPrevious() *aop.Session { @@ -433,7 +528,7 @@ type ResetSessionResponse struct { func (x *ResetSessionResponse) Reset() { *x = ResetSessionResponse{} - mi := &file_types_chat_proto_msgTypes[7] + mi := &file_types_chat_proto_msgTypes[8] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -445,7 +540,7 @@ func (x *ResetSessionResponse) String() string { func (*ResetSessionResponse) ProtoMessage() {} func (x *ResetSessionResponse) ProtoReflect() protoreflect.Message { - mi := &file_types_chat_proto_msgTypes[7] + mi := &file_types_chat_proto_msgTypes[8] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -458,7 +553,7 @@ func (x *ResetSessionResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use ResetSessionResponse.ProtoReflect.Descriptor instead. func (*ResetSessionResponse) Descriptor() ([]byte, []int) { - return file_types_chat_proto_rawDescGZIP(), []int{7} + return file_types_chat_proto_rawDescGZIP(), []int{8} } func (x *ResetSessionResponse) GetRequestId() string { @@ -519,7 +614,7 @@ type DeleteSessionRequest struct { func (x *DeleteSessionRequest) Reset() { *x = DeleteSessionRequest{} - mi := &file_types_chat_proto_msgTypes[8] + mi := &file_types_chat_proto_msgTypes[9] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -531,7 +626,7 @@ func (x *DeleteSessionRequest) String() string { func (*DeleteSessionRequest) ProtoMessage() {} func (x *DeleteSessionRequest) ProtoReflect() protoreflect.Message { - mi := &file_types_chat_proto_msgTypes[8] + mi := &file_types_chat_proto_msgTypes[9] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -544,7 +639,7 @@ func (x *DeleteSessionRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use DeleteSessionRequest.ProtoReflect.Descriptor instead. func (*DeleteSessionRequest) Descriptor() ([]byte, []int) { - return file_types_chat_proto_rawDescGZIP(), []int{8} + return file_types_chat_proto_rawDescGZIP(), []int{9} } func (x *DeleteSessionRequest) GetRequestId() string { @@ -575,7 +670,7 @@ type DeleteSessionResponse struct { func (x *DeleteSessionResponse) Reset() { *x = DeleteSessionResponse{} - mi := &file_types_chat_proto_msgTypes[9] + mi := &file_types_chat_proto_msgTypes[10] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -587,7 +682,7 @@ func (x *DeleteSessionResponse) String() string { func (*DeleteSessionResponse) ProtoMessage() {} func (x *DeleteSessionResponse) ProtoReflect() protoreflect.Message { - mi := &file_types_chat_proto_msgTypes[9] + mi := &file_types_chat_proto_msgTypes[10] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -600,7 +695,7 @@ func (x *DeleteSessionResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use DeleteSessionResponse.ProtoReflect.Descriptor instead. func (*DeleteSessionResponse) Descriptor() ([]byte, []int) { - return file_types_chat_proto_rawDescGZIP(), []int{9} + return file_types_chat_proto_rawDescGZIP(), []int{10} } func (x *DeleteSessionResponse) GetRequestId() string { @@ -660,7 +755,7 @@ type ListCommandsRequest struct { func (x *ListCommandsRequest) Reset() { *x = ListCommandsRequest{} - mi := &file_types_chat_proto_msgTypes[10] + mi := &file_types_chat_proto_msgTypes[11] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -672,7 +767,7 @@ func (x *ListCommandsRequest) String() string { func (*ListCommandsRequest) ProtoMessage() {} func (x *ListCommandsRequest) ProtoReflect() protoreflect.Message { - mi := &file_types_chat_proto_msgTypes[10] + mi := &file_types_chat_proto_msgTypes[11] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -685,7 +780,7 @@ func (x *ListCommandsRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use ListCommandsRequest.ProtoReflect.Descriptor instead. func (*ListCommandsRequest) Descriptor() ([]byte, []int) { - return file_types_chat_proto_rawDescGZIP(), []int{10} + return file_types_chat_proto_rawDescGZIP(), []int{11} } func (x *ListCommandsRequest) GetSessionId() string { @@ -704,7 +799,7 @@ type ListCommandsResponse struct { func (x *ListCommandsResponse) Reset() { *x = ListCommandsResponse{} - mi := &file_types_chat_proto_msgTypes[11] + mi := &file_types_chat_proto_msgTypes[12] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -716,7 +811,7 @@ func (x *ListCommandsResponse) String() string { func (*ListCommandsResponse) ProtoMessage() {} func (x *ListCommandsResponse) ProtoReflect() protoreflect.Message { - mi := &file_types_chat_proto_msgTypes[11] + mi := &file_types_chat_proto_msgTypes[12] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -729,7 +824,7 @@ func (x *ListCommandsResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use ListCommandsResponse.ProtoReflect.Descriptor instead. func (*ListCommandsResponse) Descriptor() ([]byte, []int) { - return file_types_chat_proto_rawDescGZIP(), []int{11} + return file_types_chat_proto_rawDescGZIP(), []int{12} } func (x *ListCommandsResponse) GetCommands() []*CommandSpec { @@ -743,7 +838,13 @@ var File_types_chat_proto protoreflect.FileDescriptor const file_types_chat_proto_rawDesc = "" + "\n" + - "\x10types/chat.proto\x12\vaiscan.chat\x1a\x0eaop/chat.proto\x1a\x13types/command.proto\x1a\x1fgoogle/protobuf/timestamp.proto\"\xe7\x01\n" + + "\x10types/chat.proto\x12\vaiscan.chat\x1a\x0eaop/chat.proto\x1a\x13types/command.proto\x1a\x1fgoogle/protobuf/timestamp.proto\"\x89\x01\n" + + "\x0eSessionHistory\x124\n" + + "\x04mode\x18\x01 \x01(\x0e2 .aiscan.chat.SessionHistory.ModeR\x04mode\"A\n" + + "\x04Mode\x12\x14\n" + + "\x10MODE_UNSPECIFIED\x10\x00\x12\x10\n" + + "\fMODE_INHERIT\x10\x01\x12\x11\n" + + "\rMODE_SNAPSHOT\x10\x02\"\xe7\x01\n" + "\rSessionRecord\x12&\n" + "\asession\x18\x01 \x01(\v2\f.aop.SessionR\asession\x12\x1d\n" + "\n" + @@ -811,43 +912,47 @@ func file_types_chat_proto_rawDescGZIP() []byte { return file_types_chat_proto_rawDescData } -var file_types_chat_proto_msgTypes = make([]protoimpl.MessageInfo, 12) +var file_types_chat_proto_enumTypes = make([]protoimpl.EnumInfo, 1) +var file_types_chat_proto_msgTypes = make([]protoimpl.MessageInfo, 13) var file_types_chat_proto_goTypes = []any{ - (*SessionRecord)(nil), // 0: aiscan.chat.SessionRecord - (*ListSessionsRequest)(nil), // 1: aiscan.chat.ListSessionsRequest - (*ListSessionsResponse)(nil), // 2: aiscan.chat.ListSessionsResponse - (*GetSessionRequest)(nil), // 3: aiscan.chat.GetSessionRequest - (*GetSessionResponse)(nil), // 4: aiscan.chat.GetSessionResponse - (*ResetSessionRequest)(nil), // 5: aiscan.chat.ResetSessionRequest - (*ResetSessionReceipt)(nil), // 6: aiscan.chat.ResetSessionReceipt - (*ResetSessionResponse)(nil), // 7: aiscan.chat.ResetSessionResponse - (*DeleteSessionRequest)(nil), // 8: aiscan.chat.DeleteSessionRequest - (*DeleteSessionResponse)(nil), // 9: aiscan.chat.DeleteSessionResponse - (*ListCommandsRequest)(nil), // 10: aiscan.chat.ListCommandsRequest - (*ListCommandsResponse)(nil), // 11: aiscan.chat.ListCommandsResponse - (*aop.Session)(nil), // 12: aop.Session - (*timestamppb.Timestamp)(nil), // 13: google.protobuf.Timestamp - (*aop.Rejection)(nil), // 14: aop.Rejection - (*CommandSpec)(nil), // 15: aiscan.command.CommandSpec + (SessionHistory_Mode)(0), // 0: aiscan.chat.SessionHistory.Mode + (*SessionHistory)(nil), // 1: aiscan.chat.SessionHistory + (*SessionRecord)(nil), // 2: aiscan.chat.SessionRecord + (*ListSessionsRequest)(nil), // 3: aiscan.chat.ListSessionsRequest + (*ListSessionsResponse)(nil), // 4: aiscan.chat.ListSessionsResponse + (*GetSessionRequest)(nil), // 5: aiscan.chat.GetSessionRequest + (*GetSessionResponse)(nil), // 6: aiscan.chat.GetSessionResponse + (*ResetSessionRequest)(nil), // 7: aiscan.chat.ResetSessionRequest + (*ResetSessionReceipt)(nil), // 8: aiscan.chat.ResetSessionReceipt + (*ResetSessionResponse)(nil), // 9: aiscan.chat.ResetSessionResponse + (*DeleteSessionRequest)(nil), // 10: aiscan.chat.DeleteSessionRequest + (*DeleteSessionResponse)(nil), // 11: aiscan.chat.DeleteSessionResponse + (*ListCommandsRequest)(nil), // 12: aiscan.chat.ListCommandsRequest + (*ListCommandsResponse)(nil), // 13: aiscan.chat.ListCommandsResponse + (*aop.Session)(nil), // 14: aop.Session + (*timestamppb.Timestamp)(nil), // 15: google.protobuf.Timestamp + (*aop.Rejection)(nil), // 16: aop.Rejection + (*CommandSpec)(nil), // 17: aiscan.command.CommandSpec } var file_types_chat_proto_depIdxs = []int32{ - 12, // 0: aiscan.chat.SessionRecord.session:type_name -> aop.Session - 13, // 1: aiscan.chat.SessionRecord.created_at:type_name -> google.protobuf.Timestamp - 13, // 2: aiscan.chat.SessionRecord.updated_at:type_name -> google.protobuf.Timestamp - 0, // 3: aiscan.chat.ListSessionsResponse.sessions:type_name -> aiscan.chat.SessionRecord - 0, // 4: aiscan.chat.GetSessionResponse.session:type_name -> aiscan.chat.SessionRecord - 12, // 5: aiscan.chat.ResetSessionReceipt.previous:type_name -> aop.Session - 0, // 6: aiscan.chat.ResetSessionReceipt.current:type_name -> aiscan.chat.SessionRecord - 6, // 7: aiscan.chat.ResetSessionResponse.accepted:type_name -> aiscan.chat.ResetSessionReceipt - 14, // 8: aiscan.chat.ResetSessionResponse.rejected:type_name -> aop.Rejection - 12, // 9: aiscan.chat.DeleteSessionResponse.accepted:type_name -> aop.Session - 14, // 10: aiscan.chat.DeleteSessionResponse.rejected:type_name -> aop.Rejection - 15, // 11: aiscan.chat.ListCommandsResponse.commands:type_name -> aiscan.command.CommandSpec - 12, // [12:12] is the sub-list for method output_type - 12, // [12:12] is the sub-list for method input_type - 12, // [12:12] is the sub-list for extension type_name - 12, // [12:12] is the sub-list for extension extendee - 0, // [0:12] is the sub-list for field type_name + 0, // 0: aiscan.chat.SessionHistory.mode:type_name -> aiscan.chat.SessionHistory.Mode + 14, // 1: aiscan.chat.SessionRecord.session:type_name -> aop.Session + 15, // 2: aiscan.chat.SessionRecord.created_at:type_name -> google.protobuf.Timestamp + 15, // 3: aiscan.chat.SessionRecord.updated_at:type_name -> google.protobuf.Timestamp + 2, // 4: aiscan.chat.ListSessionsResponse.sessions:type_name -> aiscan.chat.SessionRecord + 2, // 5: aiscan.chat.GetSessionResponse.session:type_name -> aiscan.chat.SessionRecord + 14, // 6: aiscan.chat.ResetSessionReceipt.previous:type_name -> aop.Session + 2, // 7: aiscan.chat.ResetSessionReceipt.current:type_name -> aiscan.chat.SessionRecord + 8, // 8: aiscan.chat.ResetSessionResponse.accepted:type_name -> aiscan.chat.ResetSessionReceipt + 16, // 9: aiscan.chat.ResetSessionResponse.rejected:type_name -> aop.Rejection + 14, // 10: aiscan.chat.DeleteSessionResponse.accepted:type_name -> aop.Session + 16, // 11: aiscan.chat.DeleteSessionResponse.rejected:type_name -> aop.Rejection + 17, // 12: aiscan.chat.ListCommandsResponse.commands:type_name -> aiscan.command.CommandSpec + 13, // [13:13] is the sub-list for method output_type + 13, // [13:13] is the sub-list for method input_type + 13, // [13:13] is the sub-list for extension type_name + 13, // [13:13] is the sub-list for extension extendee + 0, // [0:13] is the sub-list for field type_name } func init() { file_types_chat_proto_init() } @@ -856,11 +961,11 @@ func file_types_chat_proto_init() { return } file_types_command_proto_init() - file_types_chat_proto_msgTypes[7].OneofWrappers = []any{ + file_types_chat_proto_msgTypes[8].OneofWrappers = []any{ (*ResetSessionResponse_Accepted)(nil), (*ResetSessionResponse_Rejected)(nil), } - file_types_chat_proto_msgTypes[9].OneofWrappers = []any{ + file_types_chat_proto_msgTypes[10].OneofWrappers = []any{ (*DeleteSessionResponse_Accepted)(nil), (*DeleteSessionResponse_Rejected)(nil), } @@ -869,13 +974,14 @@ func file_types_chat_proto_init() { File: protoimpl.DescBuilder{ GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_types_chat_proto_rawDesc), len(file_types_chat_proto_rawDesc)), - NumEnums: 0, - NumMessages: 12, + NumEnums: 1, + NumMessages: 13, NumExtensions: 0, NumServices: 0, }, GoTypes: file_types_chat_proto_goTypes, DependencyIndexes: file_types_chat_proto_depIdxs, + EnumInfos: file_types_chat_proto_enumTypes, MessageInfos: file_types_chat_proto_msgTypes, }.Build() File_types_chat_proto = out.File diff --git a/pkg/types/extensions.go b/pkg/types/extensions.go index 2ae8776c..7b1232b8 100644 --- a/pkg/types/extensions.go +++ b/pkg/types/extensions.go @@ -35,6 +35,16 @@ func SetCommandDetail(event *aop.Event, value *CommandDetail) error { return aop.SetTypedExtension(event, value) } +func GetSessionHistory(event *aop.Event) (*SessionHistory, bool, error) { + value := new(SessionHistory) + ok, err := aop.FindTypedExtension(event, value) + return value, ok, err +} + +func SetSessionHistory(event *aop.Event, value *SessionHistory) error { + return aop.SetTypedExtension(event, value) +} + func GetCompactDetail(event *aop.Event) (*CompactDetail, bool, error) { value := new(CompactDetail) ok, err := aop.FindTypedExtension(event, value) diff --git a/pkg/web/service/agents_test.go b/pkg/web/service/agents_test.go index 5bfb87ff..fcf66c28 100644 --- a/pkg/web/service/agents_test.go +++ b/pkg/web/service/agents_test.go @@ -381,6 +381,25 @@ func TestWSDispatchAndComplete(t *testing.T) { } } +func TestDispatchToolCallPublishesSessionCallOnce(t *testing.T) { + sink := &evalSink{sid: "session-1", found: true} + pool := NewAgentPool(NewHub()) + pool.SetSessionLookup(sink) + remote := &remoteAgent{ + nodeState: newNodeState(), nodeID: "agent-1", + sendCh: make(chan *aop.Envelope, 1), done: make(chan struct{}), + } + pool.register(remote) + defer close(remote.done) + + if _, err := pool.DispatchToolCall("agent-1", "task-1", &aop.ToolCall{Name: "bash"}); err != nil { + t.Fatal(err) + } + if len(sink.aopEvents) != 1 || sink.aopEvents[0].GetToolCall() == nil { + t.Fatalf("session tool.call events = %+v, want exactly one hub-owned call", sink.aopEvents) + } +} + func TestWSDispatchChatUsesAOPMessage(t *testing.T) { srv, pool := setupTestServer(t) conn := dialAgentWithIdentity(t, srv, "chat-worker", []string{"scan"}, "node-chat-worker", diff --git a/pkg/web/service/broker_test.go b/pkg/web/service/broker_test.go index 46304840..3b898a00 100644 --- a/pkg/web/service/broker_test.go +++ b/pkg/web/service/broker_test.go @@ -119,6 +119,39 @@ func TestBroadcastAOPEventPersistsCanonicalProtoJSON(t *testing.T) { } } +func TestBroadcastAOPEventDoesNotFanOutRetryWithSameEventID(t *testing.T) { + store, err := NewSQLiteStore(filepath.Join(t.TempDir(), "web.db")) + if err != nil { + t.Fatal(err) + } + defer store.Close() + service := NewService(ServiceConfig{Store: store}) + createStoredSession(t, store, "session-retry") + deliveries, unsubscribe := service.SubscribeSessionEvents("session-retry") + defer unsubscribe() + event := &aop.Event{ + Id: "event-retry", SessionId: "session-retry", Emitter: "aiscan", + Payload: &aop.Event_Message{Message: &aop.Message{Id: "message-1", Role: "assistant", Content: []*aop.Content{aop.Text("once")}}}, + } + service.BroadcastAOPEvent("session-retry", event) + service.BroadcastAOPEvent("session-retry", proto.Clone(event).(*aop.Event)) + + select { + case <-deliveries: + case <-time.After(time.Second): + t.Fatal("first event was not broadcast") + } + select { + case duplicate := <-deliveries: + t.Fatalf("duplicate event was broadcast: %+v", duplicate) + case <-time.After(50 * time.Millisecond): + } + stored, err := store.ListAOPEvents(context.Background(), "session-retry", 10) + if err != nil || len(stored) != 1 { + t.Fatalf("stored events = %d, err = %v; want 1", len(stored), err) + } +} + func TestEvalMetadataPersistsOnlyInAOP(t *testing.T) { store, err := NewSQLiteStore(filepath.Join(t.TempDir(), "web.db")) if err != nil { diff --git a/pkg/web/service/events.go b/pkg/web/service/events.go index 7949ce9d..fa2ebc24 100644 --- a/pkg/web/service/events.go +++ b/pkg/web/service/events.go @@ -21,10 +21,16 @@ func (s *Service) BroadcastAOPEvent(sessionID string, event *aop.Event) { } var cursor int64 if s.store != nil { - storedCursor, _, err := s.store.AppendAOPEvent(context.Background(), sessionID, event) + storedCursor, persisted, err := s.store.AppendAOPEvent(context.Background(), sessionID, event) if err != nil { return } + // A positive cursor with persisted=false means this event ID was already + // accepted. Do not fan it out a second time. Streaming deltas are not + // persisted and intentionally have cursor 0, so they remain live-only. + if storedCursor > 0 && !persisted { + return + } cursor = storedCursor } s.broadcastAOPEvent(sessionID, event, cursor) diff --git a/pkg/web/service/store_models.go b/pkg/web/service/store_models.go index 391cf877..f30b2055 100644 --- a/pkg/web/service/store_models.go +++ b/pkg/web/service/store_models.go @@ -42,6 +42,7 @@ type aopEventModel struct { ID string `bun:"id,pk"` SessionID string `bun:"session_id,notnull,unique:aop_event_cursor"` + EventID string `bun:"event_id,notnull"` Cursor int64 `bun:"cursor,notnull,unique:aop_event_cursor"` TurnID string `bun:"turn_id,notnull"` Emitter string `bun:"emitter,notnull"` diff --git a/pkg/web/service/store_sqlite.go b/pkg/web/service/store_sqlite.go index 2d15745b..8da34ece 100644 --- a/pkg/web/service/store_sqlite.go +++ b/pkg/web/service/store_sqlite.go @@ -24,18 +24,10 @@ import ( type SQLiteStore struct { db *sql.DB orm *bun.DB - // maxEventsPerSession caps the durable AOP event rows one chat session may - // accumulate. Real sessions stay in the thousands; the cap is a capacity - // guard so an event storm (a retry loop once wrote 7.4M rows / 3.7GB into - // one session) cannot grow the database without bound. When the cap trips, - // one synthetic error event is persisted in its place so history replay - // shows why the timeline stops; live broadcast is unaffected. - maxEventsPerSession int64 } -const sqliteSchemaVersion = 3 - -const DefaultMaxAOPEventsPerSession int64 = 100_000 +// The shipped schema is a single canonical layout. Version drift is an error. +const sqliteSchemaVersion = 1 var ( dbJSONMarshal = protojson.MarshalOptions{UseProtoNames: true} @@ -50,9 +42,9 @@ func NewSQLiteStore(dbPath string) (*SQLiteStore, error) { db.SetMaxOpenConns(1) db.SetMaxIdleConns(1) orm := bun.NewDB(db, sqlitedialect.New()) - if err := migrate(orm, db); err != nil { + if err := initializeSchemaV1(orm, db); err != nil { _ = orm.Close() - return nil, fmt.Errorf("migrate sqlite: %w", err) + return nil, fmt.Errorf("initialize sqlite schema v1: %w", err) } var foreignKeys int if err := db.QueryRow(`PRAGMA foreign_keys`).Scan(&foreignKeys); err != nil || foreignKeys != 1 { @@ -62,10 +54,12 @@ func NewSQLiteStore(dbPath string) (*SQLiteStore, error) { } return nil, fmt.Errorf("verify sqlite foreign keys: disabled") } - return &SQLiteStore{db: db, orm: orm, maxEventsPerSession: DefaultMaxAOPEventsPerSession}, nil + return &SQLiteStore{db: db, orm: orm}, nil } -func migrate(orm *bun.DB, db *sql.DB) error { +// initializeSchemaV1 creates the only supported schema for a brand-new empty +// database. It never upgrades or repairs an existing database. +func initializeSchemaV1(orm *bun.DB, db *sql.DB) error { var version int if err := db.QueryRow(`PRAGMA user_version`).Scan(&version); err != nil { return err @@ -74,7 +68,7 @@ func migrate(orm *bun.DB, db *sql.DB) error { return nil } if version != 0 { - return fmt.Errorf("unsupported sqlite schema version %d; delete the database and restart", version) + return fmt.Errorf("unsupported sqlite schema version %d; database must be recreated with canonical schema v1", version) } var tables int if err := db.QueryRow(`SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name NOT LIKE 'sqlite_%'`).Scan(&tables); err != nil { @@ -126,6 +120,9 @@ func migrate(orm *bun.DB, db *sql.DB) error { return err } } + if _, err := tx.ExecContext(ctx, `CREATE UNIQUE INDEX idx_aop_events_event_id ON chat_aop_events(session_id, event_id)`); err != nil { + return err + } _, err := tx.ExecContext(ctx, fmt.Sprintf(`PRAGMA user_version = %d`, sqliteSchemaVersion)) return err }) @@ -200,10 +197,10 @@ func (s *SQLiteStore) Create(ctx context.Context, scan *types.Scan) error { func (s *SQLiteStore) Get(ctx context.Context, id string) (*types.Scan, error) { var model scanModel - if err := s.orm.NewSelect().Model(&model).Column("scan_json").Where("id = ?", id).Limit(1).Scan(ctx); err != nil { + if err := s.orm.NewSelect().Model(&model).Column("scan_json", "report").Where("id = ?", id).Limit(1).Scan(ctx); err != nil { return nil, err } - return scanFromJSON(model.ScanJSON) + return scanFromModel(model) } func (s *SQLiteStore) List(ctx context.Context, limit int) ([]*types.Scan, error) { @@ -211,12 +208,12 @@ func (s *SQLiteStore) List(ctx context.Context, limit int) ([]*types.Scan, error limit = 50 } var models []scanModel - if err := s.orm.NewSelect().Model(&models).Column("scan_json").OrderExpr("created_at DESC").Limit(limit).Scan(ctx); err != nil { + if err := s.orm.NewSelect().Model(&models).Column("scan_json", "report").OrderExpr("created_at DESC").Limit(limit).Scan(ctx); err != nil { return nil, err } scans := make([]*types.Scan, 0, len(models)) for _, model := range models { - scan, err := scanFromJSON(model.ScanJSON) + scan, err := scanFromModel(model) if err != nil { return nil, err } @@ -270,7 +267,12 @@ func scanToModel(scan *types.Scan) (*scanModel, error) { if scan == nil { return nil, fmt.Errorf("scan is required") } - raw, err := marshalProtoJSON(scan) + // Report is already stored in its dedicated relational column. Omitting it + // from the JSON snapshot avoids writing a large completed report twice while + // scanFromModel restores it for callers of Get/List. + snapshot := protobuf.CloneOf(scan) + snapshot.Report = "" + raw, err := marshalProtoJSON(snapshot) if err != nil { return nil, err } @@ -284,6 +286,17 @@ func scanToModel(scan *types.Scan) (*scanModel, error) { }, nil } +func scanFromModel(model scanModel) (*types.Scan, error) { + scan, err := scanFromJSON(model.ScanJSON) + if err != nil { + return nil, err + } + // Report has one authoritative representation: the dedicated relational + // projection. The JSON snapshot is deliberately not consulted. + scan.Report = model.Report + return scan, nil +} + func scanFromJSON(raw string) (*types.Scan, error) { scan := new(types.Scan) if err := unmarshalProtoJSON(raw, scan, "scan"); err != nil { @@ -429,6 +442,12 @@ func (s *SQLiteStore) AppendAOPEvent(ctx context.Context, sessionID string, even if event == nil || event.GetMessageDelta() != nil || event.GetToolCallDelta() != nil { return 0, false, nil } + if strings.TrimSpace(sessionID) == "" { + return 0, false, fmt.Errorf("AOP event session_id is required") + } + if strings.TrimSpace(event.Id) == "" { + return 0, false, fmt.Errorf("AOP event id is required") + } raw, err := marshalProtoJSON(event) if err != nil { return 0, false, err @@ -437,61 +456,36 @@ func (s *SQLiteStore) AppendAOPEvent(ctx context.Context, sessionID string, even if event.GetEmittedAt() != nil { createdAt = event.GetEmittedAt().AsTime().UTC().Format(time.RFC3339Nano) } - limited := false err = s.orm.RunInTx(ctx, nil, func(ctx context.Context, tx bun.Tx) error { + var existing aopEventModel + lookupErr := tx.NewSelect().Model(&existing).Column("cursor"). + Where("session_id = ? AND event_id = ?", sessionID, event.Id).Limit(1).Scan(ctx) + if lookupErr == nil { + cursor = existing.Cursor + persisted = false + return nil + } + if !errors.Is(lookupErr, sql.ErrNoRows) { + return lookupErr + } if err := tx.NewSelect().Model((*aopEventModel)(nil)). ColumnExpr("COALESCE(MAX(cursor), 0) + 1").Where("session_id = ?", sessionID).Scan(ctx, &cursor); err != nil { return err } - if s.maxEventsPerSession > 0 && cursor > s.maxEventsPerSession { - limited = true - if cursor > s.maxEventsPerSession+1 { - // The marker already terminates this session's history; - // everything past it is dropped without another row. - return nil - } - // Exactly one row past the cap: persist a synthetic terminal - // marker in place of the dropped event. The cursor sequence is - // serialized by this transaction (single-connection store), so - // the marker is written exactly once per session. - markerJSON, err := marshalProtoJSON(sessionEventLimitMarker(sessionID, s.maxEventsPerSession)) - if err != nil { - return err - } - _, err = tx.NewInsert().Model(&aopEventModel{ - ID: generateID(), SessionID: sessionID, Cursor: cursor, - Emitter: "aiscan.web", EventJSON: markerJSON, CreatedAt: createdAt, - }).Exec(ctx) - return err - } _, err := tx.NewInsert().Model(&aopEventModel{ - ID: generateID(), SessionID: sessionID, Cursor: cursor, + ID: generateID(), SessionID: sessionID, EventID: event.Id, Cursor: cursor, TurnID: event.GetTurnId(), Emitter: event.GetEmitter(), Sequence: event.GetSeq(), EventJSON: raw, CreatedAt: createdAt, }).Exec(ctx) + if err == nil { + persisted = true + } return err }) if err != nil { return 0, false, err } - if limited { - return 0, false, nil - } - return cursor, true, nil -} - -// sessionEventLimitMarker is the synthetic AOP error event persisted as a -// session's final row when it hits the event cap. It rides the normal error -// payload so replaying clients render it without special cases. -func sessionEventLimitMarker(sessionID string, limit int64) *aop.Event { - return &aop.Event{ - Id: generateID(), SessionId: sessionID, Emitter: "aiscan.web", - EmittedAt: timestamppb.Now(), - Payload: &aop.Event_Error{Error: &aop.ProtocolError{ - Code: "session_event_limit", - Message: fmt.Sprintf("session reached the %d persisted event limit; further events are not stored", limit), - }}, - } + return cursor, persisted, nil } func (s *SQLiteStore) ListAOPEvents(ctx context.Context, sessionID string, limit int) ([]*aop.Event, error) { diff --git a/pkg/web/service/store_sqlite_test.go b/pkg/web/service/store_sqlite_test.go index df48f3c7..f61aed42 100644 --- a/pkg/web/service/store_sqlite_test.go +++ b/pkg/web/service/store_sqlite_test.go @@ -4,13 +4,14 @@ import ( "context" "database/sql" "encoding/json" - "fmt" "path/filepath" + "strings" "testing" "time" aop "github.com/chainreactors/aiscan/aop" types "github.com/chainreactors/aiscan/pkg/types" + "google.golang.org/protobuf/proto" "google.golang.org/protobuf/types/known/timestamppb" ) @@ -63,7 +64,7 @@ func TestSQLiteStoreRejectsUnversionedSchema(t *testing.T) { } func TestSQLiteStoreRejectsUnsupportedSchemaVersion(t *testing.T) { - path := filepath.Join(t.TempDir(), "v1.db") + path := filepath.Join(t.TempDir(), "v2.db") db, err := sql.Open("sqlite", path) if err != nil { t.Fatal(err) @@ -78,7 +79,7 @@ func TestSQLiteStoreRejectsUnsupportedSchemaVersion(t *testing.T) { updated_at TEXT NOT NULL ); CREATE INDEX idx_sessions_agent ON chat_sessions(agent_id); - PRAGMA user_version = 1; + PRAGMA user_version = 2; `); err != nil { _ = db.Close() t.Fatal(err) @@ -90,6 +91,27 @@ func TestSQLiteStoreRejectsUnsupportedSchemaVersion(t *testing.T) { } } +func TestSQLiteStoreRejectsHistoricalSchemaVersion(t *testing.T) { + path := filepath.Join(t.TempDir(), "historical.db") + db, err := sql.Open("sqlite", path) + if err != nil { + t.Fatal(err) + } + if _, err := db.Exec(` + CREATE TABLE historical_data (id TEXT PRIMARY KEY); + PRAGMA user_version = 3; + `); err != nil { + _ = db.Close() + t.Fatal(err) + } + _ = db.Close() + + _ = db.Close() + if _, err := NewSQLiteStore(path); err == nil || !strings.Contains(err.Error(), "unsupported sqlite schema version") { + t.Fatalf("NewSQLiteStore() error = %v, want unsupported historical schema", err) + } +} + func TestSQLiteStoreAOPMessageRoundTrip(t *testing.T) { store, err := NewSQLiteStore(filepath.Join(t.TempDir(), "messages.db")) if err != nil { @@ -155,6 +177,54 @@ func TestSQLiteStoreAOPMessageRoundTrip(t *testing.T) { } } +func TestSQLiteStoreAppendAOPEventIsIdempotentByEventID(t *testing.T) { + store, err := NewSQLiteStore(filepath.Join(t.TempDir(), "aop-idempotency.db")) + if err != nil { + t.Fatal(err) + } + defer store.Close() + createStoredSession(t, store, "s1") + event := &aop.Event{ + Id: "event-retry", SessionId: "s1", Emitter: "aiscan", + Payload: &aop.Event_Message{Message: &aop.Message{Id: "m-1", Role: "assistant", Content: []*aop.Content{aop.Text("once")}}}, + } + firstCursor, firstPersisted, err := store.AppendAOPEvent(context.Background(), "s1", event) + if err != nil || !firstPersisted || firstCursor != 1 { + t.Fatalf("first append = cursor:%d persisted:%v err:%v", firstCursor, firstPersisted, err) + } + var storedEventID string + if err := store.db.QueryRow(`SELECT event_id FROM chat_aop_events WHERE session_id = ?`, "s1").Scan(&storedEventID); err != nil { + t.Fatal(err) + } + if storedEventID != event.Id { + t.Fatalf("stored event id = %q, want %q", storedEventID, event.Id) + } + secondCursor, secondPersisted, err := store.AppendAOPEvent(context.Background(), "s1", proto.Clone(event).(*aop.Event)) + if err != nil || secondPersisted || secondCursor != firstCursor { + t.Fatalf("retry append = cursor:%d persisted:%v err:%v", secondCursor, secondPersisted, err) + } + var count int + if err := store.db.QueryRow(`SELECT COUNT(*) FROM chat_aop_events WHERE session_id = ?`, "s1").Scan(&count); err != nil { + t.Fatal(err) + } + if count != 1 { + t.Fatalf("persisted event count = %d, want 1", count) + } +} + +func TestSQLiteStoreRejectsAOPEventWithoutIdentity(t *testing.T) { + store, err := NewSQLiteStore(filepath.Join(t.TempDir(), "missing-event-id.db")) + if err != nil { + t.Fatal(err) + } + defer store.Close() + if _, _, err := store.AppendAOPEvent(context.Background(), "s1", &aop.Event{ + SessionId: "s1", Payload: &aop.Event_Status{Status: &aop.Status{State: "ready"}}, + }); err == nil { + t.Fatal("AppendAOPEvent accepted an event without an id") + } +} + func TestSQLiteStorePersistsAnalysisOptions(t *testing.T) { store, err := NewSQLiteStore(filepath.Join(t.TempDir(), "scans.db")) if err != nil { @@ -228,6 +298,42 @@ func TestSQLiteStoreUsesProtoJSONAndRelationalScanColumns(t *testing.T) { } } +func TestSQLiteStoreDoesNotDuplicateLargeReportInSnapshot(t *testing.T) { + store, err := NewSQLiteStore(filepath.Join(t.TempDir(), "report-dedup.db")) + if err != nil { + t.Fatal(err) + } + defer store.Close() + + report := strings.Repeat("report-line\n", 4<<20/12) + scan := &types.Scan{ + Id: "scan-report-dedup", Target: "example.com", Mode: "quick", Report: report, + Status: types.ScanStatus_SCAN_STATUS_COMPLETED, CreatedAt: nowProto(), UpdatedAt: nowProto(), + } + if err := store.Create(context.Background(), scan); err != nil { + t.Fatal(err) + } + var snapshotBytes, reportBytes int + var raw string + if err := store.db.QueryRow(`SELECT scan_json, length(scan_json), length(report) FROM scans WHERE id = ?`, scan.Id). + Scan(&raw, &snapshotBytes, &reportBytes); err != nil { + t.Fatal(err) + } + if reportBytes != len(report) { + t.Fatalf("report column bytes = %d, want %d", reportBytes, len(report)) + } + if strings.Contains(raw, "report-line") || snapshotBytes >= len(report) { + t.Fatalf("scan_json still duplicates the large report: snapshot_bytes=%d report_bytes=%d", snapshotBytes, reportBytes) + } + got, err := store.Get(context.Background(), scan.Id) + if err != nil { + t.Fatal(err) + } + if got.Report != report { + t.Fatalf("Get() report bytes = %d, want %d", len(got.Report), len(report)) + } +} + func TestSQLiteStoreKeepsSCOObservationForEveryOperation(t *testing.T) { store, err := NewSQLiteStore(filepath.Join(t.TempDir(), "sco.db")) if err != nil { @@ -343,52 +449,6 @@ func TestSQLiteStoreEnablesForeignKeysAndCascadesSessionData(t *testing.T) { } } -func TestSQLiteStoreCapsAOPEventsPerSession(t *testing.T) { - store, err := NewSQLiteStore(filepath.Join(t.TempDir(), "event-cap.db")) - if err != nil { - t.Fatal(err) - } - defer store.Close() - store.maxEventsPerSession = 5 - ctx := context.Background() - createStoredSession(t, store, "s1") - - makeEvent := func(i int) *aop.Event { - return &aop.Event{ - Id: fmt.Sprintf("e-%d", i), EmittedAt: timestamppb.Now(), SessionId: "s1", Emitter: "aiscan", - Payload: &aop.Event_Message{Message: &aop.Message{ - Id: fmt.Sprintf("m-%d", i), Role: "assistant", Content: []*aop.Content{aop.Text(fmt.Sprintf("event %d", i))}, - }}, - } - } - for i := 1; i <= 5; i++ { - cursor, persisted, err := store.AppendAOPEvent(ctx, "s1", makeEvent(i)) - if err != nil || !persisted || cursor != int64(i) { - t.Fatalf("event %d: cursor=%d persisted=%v err=%v", i, cursor, persisted, err) - } - } - // Everything past the cap is dropped without error so live broadcast - // keeps working; only one synthetic marker row is added. - for i := 6; i <= 10; i++ { - cursor, persisted, err := store.AppendAOPEvent(ctx, "s1", makeEvent(i)) - if err != nil || persisted || cursor != 0 { - t.Fatalf("event %d past cap: cursor=%d persisted=%v err=%v", i, cursor, persisted, err) - } - } - - events, err := store.ListAOPEvents(ctx, "s1", 100) - if err != nil { - t.Fatal(err) - } - if len(events) != 6 { - t.Fatalf("stored events = %d, want 5 real + 1 marker", len(events)) - } - marker := events[len(events)-1].GetError() - if marker == nil || marker.Code != "session_event_limit" { - t.Fatalf("final event = %+v, want session_event_limit marker", events[len(events)-1]) - } -} - func TestSQLiteStoreRejectsAOPEventForMissingSession(t *testing.T) { store, err := NewSQLiteStore(filepath.Join(t.TempDir(), "foreign-keys.db")) if err != nil { diff --git a/proto/types/chat.proto b/proto/types/chat.proto index ba401a72..c7541fac 100644 --- a/proto/types/chat.proto +++ b/proto/types/chat.proto @@ -8,6 +8,17 @@ import "google/protobuf/timestamp.proto"; option go_package = "github.com/chainreactors/aiscan/pkg/types;types"; +// SessionHistory is persisted as an AOP event extension. It makes transcript +// inheritance explicit without changing the shared AOP protocol schema. +message SessionHistory { + enum Mode { + MODE_UNSPECIFIED = 0; + MODE_INHERIT = 1; + MODE_SNAPSHOT = 2; + } + Mode mode = 1; +} + message SessionRecord { aop.Session session = 1; string agent_name = 2; diff --git a/tools/curl/client.go b/tools/curl/client.go index a89952d3..2176688a 100644 --- a/tools/curl/client.go +++ b/tools/curl/client.go @@ -26,6 +26,7 @@ import ( toolpb "github.com/chainreactors/aiscan/aop/tool" traffic "github.com/chainreactors/aiscan/aop/traffic" + "github.com/chainreactors/aiscan/pkg/commands" ) // A single stable, modern Chrome identity. Keeping one fingerprint per process @@ -53,16 +54,17 @@ var browserDefaults = []Header{ } // do runs one parsed curl request end to end: builds the client (routing through -// the runner's MITM hub when the environment provides it), applies the browser -// naturalization defaults, performs the exchange, and writes curl-shaped output. -// env and workDir are per-invocation; nothing here mutates the shared Command. -func (c *Command) do(ctx context.Context, req *Request, env map[string]string, workDir string, stdout, stderr io.Writer) error { +// the runner's MITM hub when the resolved egress provides it), applies the +// browser naturalization defaults, performs the exchange, and writes +// curl-shaped output. Egress and workDir are per-invocation; nothing here +// mutates the shared Command. +func (c *Command) do(ctx context.Context, req *Request, egress commands.Egress, workDir string, stdout, stderr io.Writer) error { if req.Version { _, err := fmt.Fprintln(stdout, compatibilityVersion) return err } - proxyURL, caPath := c.egress(env) + proxyURL, caPath := egress.ProxyURL, egress.CAPath if req.Proxy != "" { // -x overrides the injected hub egress for this invocation only. proxyURL = req.Proxy @@ -232,29 +234,6 @@ func isTimeoutError(err error) bool { return errors.As(err, &netErr) && netErr.Timeout() } -// egress reads the hub proxy and CA path the runner injected into this -// execution's environment. The proxy URL already carries the tool-call id as -// its username, so captured flows attribute to this call; the CA is present -// only while the hub is intercepting. Falls back to the static scanner proxy. -func (c *Command) egress(env map[string]string) (proxyURL, caPath string) { - for _, key := range []string{"ALL_PROXY", "all_proxy", "HTTPS_PROXY", "https_proxy", "HTTP_PROXY", "http_proxy"} { - if v := env[key]; v != "" { - proxyURL = v - break - } - } - for _, key := range []string{"CURL_CA_BUNDLE", "SSL_CERT_FILE"} { - if v := env[key]; v != "" { - caPath = v - break - } - } - if proxyURL == "" { - proxyURL = c.Proxy - } - return proxyURL, caPath -} - func (c *Command) buildClient(proxyURL, caPath string, req *Request) (*http.Client, error) { tlsConfig := &tls.Config{MinVersion: tls.VersionTLS12} switch { diff --git a/tools/curl/client_test.go b/tools/curl/client_test.go index 9876bcaa..85df63ce 100644 --- a/tools/curl/client_test.go +++ b/tools/curl/client_test.go @@ -12,6 +12,8 @@ import ( "strings" "testing" "time" + + "github.com/chainreactors/aiscan/pkg/commands" ) // run is a small harness: parse args, execute against a real server, capture @@ -24,13 +26,12 @@ func run(t *testing.T, args []string, env, workDir string) (stdout, stderr strin return "", "", perr } var out, errb strings.Builder - envMapVal := map[string]string{} + c := New() + var egress commands.Egress if env != "" { - k, v, _ := strings.Cut(env, "=") - envMapVal[k] = v + egress = commands.ResolveEgress([]string{env}, c.Proxy) } - c := New() - err = c.do(context.Background(), req, envMapVal, workDir, &out, &errb) + err = c.do(context.Background(), req, egress, workDir, &out, &errb) return out.String(), errb.String(), err } @@ -330,6 +331,28 @@ func TestProxyOverride(t *testing.T) { } } +func TestRunnerEgressProxy(t *testing.T) { + proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.HasPrefix(r.RequestURI, "http://") { + t.Errorf("proxy got non-absolute URI %q", r.RequestURI) + } + _, _ = w.Write([]byte("via-runner-egress")) + })) + defer proxy.Close() + target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte("direct")) + })) + defer target.Close() + + out, _, err := run(t, []string{target.URL}, "ALL_PROXY="+proxy.URL, "") + if err != nil { + t.Fatal(err) + } + if out != "via-runner-egress" { + t.Fatalf("body = %q, want via-runner-egress", out) + } +} + func TestDumpHeaders(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("X-Dump-Test", "yes") diff --git a/tools/curl/curl.go b/tools/curl/curl.go index 9286a437..c0560285 100644 --- a/tools/curl/curl.go +++ b/tools/curl/curl.go @@ -3,7 +3,6 @@ package curl import ( "context" "fmt" - "strings" aop "github.com/chainreactors/aiscan/aop" "github.com/chainreactors/aiscan/core/telemetry" @@ -126,15 +125,5 @@ func (c *Command) Run(ctx context.Context, execution *commands.Execution) (_ any if workDir == "" { workDir = coretool.WorkDirFromContext(ctx, c.WorkDir) } - return nil, c.do(ctx, req, envMap(execution.Env), workDir, execution.Stdout, execution.Stderr) -} - -func envMap(env []string) map[string]string { - values := make(map[string]string, len(env)) - for _, item := range env { - if key, value, ok := strings.Cut(item, "="); ok { - values[key] = value - } - } - return values + return nil, c.do(ctx, req, commands.ResolveExecutionEgress(execution, c.Proxy), workDir, execution.Stdout, execution.Stderr) } diff --git a/tools/gogo/gogo.go b/tools/gogo/gogo.go index 7820e997..a877a80c 100644 --- a/tools/gogo/gogo.go +++ b/tools/gogo/gogo.go @@ -68,7 +68,8 @@ func (c *Command) Run(ctx context.Context, execution *commands.Execution) (_ any defer telemetry.RecoverAsError("gogo", &err) args := execution.Args args = c.normalizeArgs(args) - args = c.injectProxy(args) + egress := commands.ResolveExecutionEgress(execution, c.Proxy) + args = c.injectProxyURL(args, egress.ProxyURL) if toolargs.BoolFlagEnabled(args, "--debug") { restoreDebug := telemetry.ActivateDebug(c.Logger) @@ -109,13 +110,17 @@ func (c *Command) TestInjectProxy(args []string) []string { } func (c *Command) injectProxy(args []string) []string { - if c.Proxy == "" { + return c.injectProxyURL(args, c.Proxy) +} + +func (c *Command) injectProxyURL(args []string, proxy string) []string { + if proxy == "" { return args } if toolargs.HasFlag(args, "--proxy") { return args } - return append(args, "--proxy", c.Proxy) + return append(args, "--proxy", proxy) } // normalizeArgs adapts common agent-generated gogo arguments before handing diff --git a/tools/katana/katana.go b/tools/katana/katana.go index c72ba092..8367c4cd 100644 --- a/tools/katana/katana.go +++ b/tools/katana/katana.go @@ -145,9 +145,11 @@ func (c *Command) Run(ctx context.Context, execution *commands.Execution) (_ any options.NoColors = true options.DisableUpdateCheck = true - // Inject proxy. - if options.Proxy == "" && c.Proxy != "" { - options.Proxy = c.Proxy + // Inject the call-scoped Runner route only when the caller did not provide + // an explicit -proxy option. + egress := commands.ResolveExecutionEgress(execution, c.Proxy) + if options.Proxy == "" { + options.Proxy = egress.ProxyURL } if err := configureBrowserOptions(options); err != nil { return nil, fmt.Errorf("katana: %w", err) diff --git a/tools/neutron/neutron.go b/tools/neutron/neutron.go index eda55e3d..a9a7e47b 100644 --- a/tools/neutron/neutron.go +++ b/tools/neutron/neutron.go @@ -159,7 +159,12 @@ func (c *Command) Run(ctx context.Context, execution *commands.Execution) (_ any return nil, fmt.Errorf("neutron: --rate-limit cannot be negative") } - loadedTemplates, err := loadNeutronTemplatePaths(flags.Templates) + // The Runner injects a call-scoped Hub URL into Execution.Env. Resolve it + // before loading -t templates: the SDK binds each request transport at + // compile time, so a zero-proxy loader would let an explicit template dial + // the target directly even though the command itself has a proxy configured. + proxyURL := commands.ResolveExecutionEgress(execution, c.Proxy).ProxyURL + loadedTemplates, err := loadNeutronTemplatePaths(flags.Templates, proxyURL) if err != nil { return nil, err } @@ -336,11 +341,11 @@ func readNeutronTargets(inputs []string, input, listFile string) ([]string, erro return out, scanner.Err() } -func loadNeutronTemplatePaths(paths []string) ([]*templates.Template, error) { +func loadNeutronTemplatePaths(paths []string, proxyURL string) ([]*templates.Template, error) { if len(paths) == 0 { return nil, nil } - cfg := sdkneutron.NewConfig() + cfg := sdkneutron.NewConfig().WithProxy(proxyURL) engine, err := sdkneutron.NewEngine(cfg.WithTemplates([]*templates.Template{minimalCompilableTemplate()})) if err != nil { return nil, fmt.Errorf("neutron: initialize template loader: %w", err) diff --git a/tools/neutron/neutron_test.go b/tools/neutron/neutron_test.go index 43ab56a5..73e2d035 100644 --- a/tools/neutron/neutron_test.go +++ b/tools/neutron/neutron_test.go @@ -4,10 +4,13 @@ import ( "bytes" "context" "encoding/json" + "net" + nethttp "net/http" "os" "path/filepath" "strings" "testing" + "time" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/neutron/operators" @@ -127,6 +130,67 @@ http: } } +func TestExplicitTemplateLoaderBindsCallProxy(t *testing.T) { + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer listener.Close() + hit := make(chan struct{}, 1) + go func() { + conn, acceptErr := listener.Accept() + if acceptErr == nil { + select { + case hit <- struct{}{}: + default: + } + _ = conn.Close() + } + }() + + path := filepath.Join(t.TempDir(), "custom.yml") + if err := os.WriteFile(path, []byte(`id: proxy-bound +info: + name: proxy-bound + severity: info +http: + - method: GET + path: + - '{{BaseURL}}' +`), 0600); err != nil { + t.Fatal(err) + } + + proxyURL := "http://" + listener.Addr().String() + loaded, err := loadNeutronTemplatePaths([]string{path}, proxyURL) + if err != nil { + t.Fatal(err) + } + if len(loaded) != 1 || len(loaded[0].GetRequests()) != 1 { + t.Fatalf("loaded templates = %#v, want one HTTP template", loaded) + } + client := loaded[0].GetRequests()[0].GetHTTPClient() + if client == nil { + t.Fatal("compiled request has no HTTP client") + } + transport, ok := client.Transport.(*nethttp.Transport) + if !ok || transport.DialContext == nil { + t.Fatalf("compiled transport = %#v, want proxy dialer", client.Transport) + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + conn, _ := transport.DialContext(ctx, "tcp", "example.invalid:80") + if conn != nil { + _ = conn.Close() + } + select { + case <-hit: + case <-time.After(time.Second): + t.Fatal("compiled template did not dial the configured proxy") + } +} + func newTestNeutronEngine(t *testing.T, items ...*templates.Template) *sdkneutron.Engine { t.Helper() engine, err := sdkneutron.NewEngineWithTemplates((sdkneutron.Templates{}).Merge(items)) diff --git a/tools/scan/adapter.go b/tools/scan/adapter.go index 812a6a20..79171a27 100644 --- a/tools/scan/adapter.go +++ b/tools/scan/adapter.go @@ -34,7 +34,7 @@ func (c *Command) runPortDiscoveryCapability(ctx context.Context, discovery disc Timeout: discovery.Timeout, VersionLevel: discovery.Version, Exploit: discovery.Exploit, - Proxy: c.Proxy, + Proxy: c.proxyForContext(ctx), Debug: discovery.Debug, OnStats: func(stats sdktypes.Stats) { emit(statsEvent(capGogoPortscan, stats)) @@ -62,6 +62,7 @@ func (c *Command) runSprayCapability(ctx context.Context, flags flags, web webOp return } opts = applyWebStrategyOptions(flags, web, opts) + opts.Proxy = c.proxyForContext(ctx) opts.URLs = []string{target.URL} opts.Host = target.HostHeader opts.Scope = webTargetScope(target) @@ -123,7 +124,7 @@ func (c *Command) runWeakpassCapability(ctx context.Context, flags flags, creden Top: flags.ZombieTop, Users: credentials.Users, Passwords: credentials.Passwords, - Proxy: c.Proxy, + Proxy: c.proxyForContext(ctx), Debug: flags.Debug, OnStats: func(stats sdktypes.Stats) { emit(statsEvent(capZombieWeakpass, stats)) @@ -221,7 +222,7 @@ func (c *Command) runHTTPBasicAuthCapability(ctx context.Context, flags flags, i if !ok || !reportableSprayResultForCapability(target.Result, target.Capability) || target.Result.Status != 401 { return } - zTarget, ok := basicAuthZombieTarget(ctx, target.Result.UrlString, target.HostHeader, flags.Timeout) + zTarget, ok := basicAuthZombieTarget(ctx, target.Result.UrlString, target.HostHeader, flags.Timeout, c.proxyForContext(ctx)) if !ok { return } diff --git a/tools/scan/capability.go b/tools/scan/capability.go index f18f54e1..69f18a9b 100644 --- a/tools/scan/capability.go +++ b/tools/scan/capability.go @@ -96,6 +96,9 @@ func (c *Command) buildCapabilities(flags flags, opts scanOptions, profile profi if !profile.Enabled(name) || !hasSpray(c.engines) { return } + // The final call-scoped route is applied in runSprayCapability, where the + // pipeline context is available. Keep the startup value here for direct + // unit callers that do not install an invocation context. sopts.Proxy = c.Proxy sprayBuilt = true capabilities = append(capabilities, sprayCapability(c, flags, opts.Web, name, sources, sopts, c.runSprayCapability)) @@ -129,7 +132,7 @@ func (c *Command) buildCapabilities(flags flags, opts scanOptions, profile profi wrapRoutes(acceptsTarget(targetWeb), webSources()...), capWorkers(c.engines.Capacity.Spray, flags.SprayThreads), func(ctx context.Context, e event, emit func(event)) { - c.runSprayCapability(ctx, flags, opts.Web, e.Target, capSprayCrawl, engine.SprayCheckOptions{Crawl: true, CrawlDepth: profile.CrawlDepth, Proxy: c.Proxy}, emit) + c.runSprayCapability(ctx, flags, opts.Web, e.Target, capSprayCrawl, engine.SprayCheckOptions{Crawl: true, CrawlDepth: profile.CrawlDepth, Proxy: c.proxyForContext(ctx)}, emit) }, )) } diff --git a/tools/scan/capability_katana.go b/tools/scan/capability_katana.go index e8061ef7..82f8fc1b 100644 --- a/tools/scan/capability_katana.go +++ b/tools/scan/capability_katana.go @@ -135,8 +135,8 @@ func runKatanaCrawl(ctx context.Context, c *Command, e event, depth int, jsMode handleResult(&r) }, } - if c.Proxy != "" { - options.Proxy = c.Proxy + if proxy := c.proxyForContext(ctx); proxy != "" { + options.Proxy = proxy } if jsMode { binary, err := browserutil.Discover() diff --git a/tools/scan/command.go b/tools/scan/command.go index c0ceeff1..3582744b 100644 --- a/tools/scan/command.go +++ b/tools/scan/command.go @@ -85,6 +85,8 @@ func Usage() string { func (c *Command) Run(ctx context.Context, execution *commands.Execution) (_ any, err error) { defer telemetry.RecoverAsError("scan", &err) + egress := commands.ResolveExecutionEgress(execution, c.Proxy) + ctx = withInvocationProxy(ctx, egress.ProxyURL) out, _, err := c.execute(ctx, c.resolveRelativePaths(execution.Args), execution.Stdout) if err != nil { return nil, err diff --git a/tools/scan/http_auth.go b/tools/scan/http_auth.go index 567ca7b3..19a45ffc 100644 --- a/tools/scan/http_auth.go +++ b/tools/scan/http_auth.go @@ -3,6 +3,7 @@ package scan import ( "context" "crypto/tls" + "fmt" "net/http" "net/url" "regexp" @@ -15,12 +16,12 @@ import ( var basicAuthChallengePattern = regexp.MustCompile(`(?i)(^|,)\s*basic(\s|$)`) -func basicAuthZombieTarget(ctx context.Context, rawURL, hostHeader string, timeoutSeconds int) (sdkzombie.Target, bool) { +func basicAuthZombieTarget(ctx context.Context, rawURL, hostHeader string, timeoutSeconds int, proxy string) (sdkzombie.Target, bool) { parsed, ok := parseInputURL(rawURL) if !ok || !utils.IsWebScheme(parsed.Scheme) { return sdkzombie.Target{}, false } - if !hasHTTPBasicAuthChallenge(ctx, parsed, hostHeader, timeoutSeconds) { + if !hasHTTPBasicAuthChallenge(ctx, parsed, hostHeader, timeoutSeconds, proxy) { return sdkzombie.Target{}, false } @@ -32,7 +33,7 @@ func basicAuthZombieTarget(ctx context.Context, rawURL, hostHeader string, timeo return target, true } -func hasHTTPBasicAuthChallenge(ctx context.Context, parsed *url.URL, hostHeader string, timeoutSeconds int) bool { +func hasHTTPBasicAuthChallenge(ctx context.Context, parsed *url.URL, hostHeader string, timeoutSeconds int, proxy string) bool { if parsed == nil { return false } @@ -48,7 +49,7 @@ func hasHTTPBasicAuthChallenge(ctx context.Context, parsed *url.URL, hostHeader req.Header.Set("User-Agent", "aiscan") req.Close = true - client := httpAuthClient(timeoutSeconds) + client := httpAuthClient(timeoutSeconds, proxy) defer client.CloseIdleConnections() resp, err := client.Do(req) if err != nil { @@ -58,11 +59,31 @@ func hasHTTPBasicAuthChallenge(ctx context.Context, parsed *url.URL, hostHeader return resp.StatusCode == http.StatusUnauthorized && hasBasicAuthChallenge(resp.Header.Values("WWW-Authenticate")) } -func httpAuthClient(timeoutSeconds int) *http.Client { +func httpAuthClient(timeoutSeconds int, proxy string) *http.Client { if timeoutSeconds <= 0 { timeoutSeconds = 5 } - transport := http.DefaultTransport.(*http.Transport).Clone() //nolint:errcheck // DefaultTransport is always *http.Transport + transport := http.DefaultTransport.(*http.Transport).Clone() //nolint:errcheck // DefaultTransport is always *http.Transport + // Do not inherit the process environment: Runner egress is resolved from + // the invocation and an ambient proxy could bypass the Hub contract. + transport.Proxy = nil + if strings.TrimSpace(proxy) != "" { + rawProxy := strings.TrimSpace(proxy) + parsed, err := url.Parse(rawProxy) + if err != nil || parsed.Scheme == "" || parsed.Host == "" { + // Fail closed: a malformed call-scoped route must never turn this + // probe into a direct request. + proxyErr := err + if proxyErr == nil { + proxyErr = fmt.Errorf("expected URL with scheme and host") + } + transport.Proxy = func(*http.Request) (*url.URL, error) { + return nil, fmt.Errorf("invalid proxy %q: %w", rawProxy, proxyErr) + } + } else { + transport.Proxy = http.ProxyURL(parsed) + } + } transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} //nolint:gosec // scanner probes must tolerate self-signed certs return &http.Client{ Timeout: time.Duration(timeoutSeconds) * time.Second, diff --git a/tools/scan/options.go b/tools/scan/options.go index e9870e1c..e6c79b78 100644 --- a/tools/scan/options.go +++ b/tools/scan/options.go @@ -8,6 +8,25 @@ import ( "github.com/chainreactors/aiscan/core/telemetry" ) +type invocationProxyKey struct{} + +func withInvocationProxy(ctx context.Context, proxy string) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, invocationProxyKey{}, proxy) +} + +func (c *Command) proxyForContext(ctx context.Context) string { + if ctx == nil { + return c.Proxy + } + if proxy, ok := ctx.Value(invocationProxyKey{}).(string); ok { + return proxy + } + return c.Proxy +} + type Option func(*Command) type DeepBrowserFunc func(ctx context.Context, targetURL string) (string, error) diff --git a/tools/search/fetch.go b/tools/search/fetch.go index f5bdd777..7c9d38ab 100644 --- a/tools/search/fetch.go +++ b/tools/search/fetch.go @@ -2,10 +2,13 @@ package search import ( "context" + "crypto/tls" + "crypto/x509" "fmt" "io" "net/http" "net/url" + "os" "regexp" "strings" "sync" @@ -128,8 +131,9 @@ func (c *urlCache) Clear() { // --------------------------------------------------------------------------- type FetchCommand struct { - client *http.Client - cache *urlCache + cache *urlCache + proxy string + ca string } func (c *FetchCommand) Name() string { return "fetch" } @@ -142,21 +146,62 @@ Useful for reading advisories, documentation, and vulnerability details.` } func NewFetchCommand() *FetchCommand { - transport := &http.Transport{ - Proxy: http.ProxyFromEnvironment, - } return &FetchCommand{ - client: &http.Client{ - Transport: transport, - Timeout: fetchTimeout, - CheckRedirect: func(*http.Request, []*http.Request) error { - return http.ErrUseLastResponse - }, - }, cache: newURLCache(), } } +// WithProxy sets the startup fallback. Each Run still resolves the +// call-scoped Runner environment before constructing its private client. +func (c *FetchCommand) WithProxy(proxy string) *FetchCommand { + c.proxy = proxy + return c +} + +func (c *FetchCommand) WithProxyCA(ca string) *FetchCommand { + c.ca = ca + return c +} + +func fetchClientForProxy(proxy, caPath string) (*http.Client, error) { + transport := http.DefaultTransport.(*http.Transport).Clone() //nolint:errcheck // DefaultTransport is always *http.Transport + // Never inherit an ambient process proxy; the invocation's resolved route is + // the sole egress authority for this request. + transport.Proxy = nil + proxy = strings.TrimSpace(proxy) + if proxy != "" { + u, err := url.Parse(proxy) + if err != nil || u.Scheme == "" || u.Host == "" { + if err == nil { + err = fmt.Errorf("expected URL with scheme and host") + } + return nil, fmt.Errorf("fetch: invalid proxy %q: %w", proxy, err) + } + transport.Proxy = http.ProxyURL(u) + } + if caPath = strings.TrimSpace(caPath); caPath != "" { + pool, err := x509.SystemCertPool() + if err != nil || pool == nil { + pool = x509.NewCertPool() + } + pem, err := os.ReadFile(caPath) + if err != nil { + return nil, fmt.Errorf("fetch: read CA bundle %q: %w", caPath, err) + } + if !pool.AppendCertsFromPEM(pem) { + return nil, fmt.Errorf("fetch: CA bundle %q contains no certificates", caPath) + } + transport.TLSClientConfig = &tls.Config{RootCAs: pool, MinVersion: tls.VersionTLS12} + } + return &http.Client{ + Transport: transport, + Timeout: fetchTimeout, + CheckRedirect: func(*http.Request, []*http.Request) error { + return http.ErrUseLastResponse + }, + }, nil +} + func (c *FetchCommand) ClearCache() { c.cache.Clear() } func (c *FetchCommand) Run(ctx context.Context, execution *commands.Execution) (_ any, err error) { @@ -184,7 +229,12 @@ func (c *FetchCommand) Run(ctx context.Context, execution *commands.Execution) ( return nil, nil } - result, redir, err := c.fetchWithRedirects(ctx, normalizedURL, 0) + egress := commands.ResolveExecutionEgress(execution, c.proxy) + client, err := fetchClientForProxy(egress.ProxyURL, egress.CAPath) + if err != nil { + return nil, err + } + result, redir, err := c.fetchWithRedirects(ctx, client, normalizedURL, 0) if err != nil { return nil, err } @@ -247,7 +297,7 @@ type redirectInfo struct { statusCode int } -func (c *FetchCommand) fetchWithRedirects(ctx context.Context, targetURL string, depth int) (*fetchResult, *redirectInfo, error) { +func (c *FetchCommand) fetchWithRedirects(ctx context.Context, client *http.Client, targetURL string, depth int) (*fetchResult, *redirectInfo, error) { if depth > maxRedirects { return nil, nil, fmt.Errorf("too many redirects (exceeded %d)", maxRedirects) } @@ -260,7 +310,7 @@ func (c *FetchCommand) fetchWithRedirects(ctx context.Context, targetURL string, req.Header.Set("Accept", "text/markdown, text/html, text/plain, */*") req.Header.Set("Accept-Language", "en-US,en;q=0.9,zh-CN;q=0.8,zh;q=0.7") - resp, err := c.client.Do(req) + resp, err := client.Do(req) if err != nil { return nil, nil, fmt.Errorf("fetch failed: %w", err) } @@ -277,7 +327,7 @@ func (c *FetchCommand) fetchWithRedirects(ctx context.Context, targetURL string, } if isPermittedRedirect(targetURL, redirectURL) { - return c.fetchWithRedirects(ctx, redirectURL, depth+1) + return c.fetchWithRedirects(ctx, client, redirectURL, depth+1) } return nil, &redirectInfo{ originalURL: targetURL, diff --git a/tools/search/fetch_test.go b/tools/search/fetch_test.go index 3c0c012a..cf206d0f 100644 --- a/tools/search/fetch_test.go +++ b/tools/search/fetch_test.go @@ -38,6 +38,34 @@ func TestFetchExecutePreservesExplicitHTTPURL(t *testing.T) { } } +func TestFetchUsesCallScopedRunnerProxy(t *testing.T) { + var proxyHits int + proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + proxyHits++ + if r.URL.Path != "/through-proxy" { + t.Errorf("proxy request path = %q, want /through-proxy", r.URL.Path) + } + w.Header().Set("Content-Type", "text/plain") + _, _ = w.Write([]byte("proxied")) + })) + defer proxy.Close() + + cmd := NewFetchCommand() + var output bytes.Buffer + _, err := cmd.Run(context.Background(), &commands.Execution{ + Args: []string{proxy.URL + "/through-proxy"}, + Env: []string{"ALL_PROXY=" + proxy.URL}, + Stdout: &output, + Stderr: &output, + }) + if err != nil { + t.Fatalf("fetch through call-scoped proxy: %v", err) + } + if proxyHits != 1 || !strings.Contains(output.String(), "proxied") { + t.Fatalf("proxy hits=%d output=%q, want one proxied response", proxyHits, output.String()) + } +} + func TestFetchCacheHitReturnsCachedContent(t *testing.T) { callCount := 0 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { diff --git a/tools/search/register.go b/tools/search/register.go index 8995dc1a..6d41b011 100644 --- a/tools/search/register.go +++ b/tools/search/register.go @@ -23,7 +23,7 @@ func init() { } reg.RegisterTool(NewWebSearchTool(d.Provider, tavily)) - fetch := NewFetchCommand() + fetch := NewFetchCommand().WithProxy(d.ScannerProxy).WithProxyCA(d.ScannerProxyCA) reg.Register(commands.Command{ Name: fetch.Name(), Usage: fetch.Usage(), DescriptionPath: "aiscan://skills/aiscan/okf/runtime/fetch.md", diff --git a/tools/search/tavily.go b/tools/search/tavily.go index a339a6e0..8da95203 100644 --- a/tools/search/tavily.go +++ b/tools/search/tavily.go @@ -79,16 +79,27 @@ func NewTavilySearch(builtinKeys string) *TavilySearch { return c } -// proxyTransport builds an HTTP transport that routes through proxy when it is a -// valid URL, falling back to the environment proxy for an empty or unparseable -// value. +// proxyTransport builds an HTTP transport from the configured egress. It never +// consults the process environment: an invocation's route must not be replaced +// by ambient proxy state, and malformed routes fail closed in Transport.Proxy. func proxyTransport(proxy string) *http.Transport { - t := &http.Transport{Proxy: http.ProxyFromEnvironment} - if proxy != "" { - if u, err := url.Parse(proxy); err == nil { - t.Proxy = http.ProxyURL(u) + t := &http.Transport{} + proxy = strings.TrimSpace(proxy) + if proxy == "" { + return t + } + u, err := url.Parse(proxy) + if err != nil || u.Scheme == "" || u.Host == "" { + proxyErr := err + if proxyErr == nil { + proxyErr = fmt.Errorf("expected URL with scheme and host") } + t.Proxy = func(*http.Request) (*url.URL, error) { + return nil, fmt.Errorf("invalid proxy %q: %w", proxy, proxyErr) + } + return t } + t.Proxy = http.ProxyURL(u) return t } diff --git a/tools/spray/spray.go b/tools/spray/spray.go index 85e4910b..f64cbfab 100644 --- a/tools/spray/spray.go +++ b/tools/spray/spray.go @@ -80,7 +80,8 @@ func (c *Command) Run(ctx context.Context, execution *commands.Execution) (_ any if c.engine != nil { c.engine.InstallResourceProvider() } - args = c.injectProxy(args) + egress := commands.ResolveExecutionEgress(execution, c.Proxy) + args = c.injectProxyURL(args, egress.ProxyURL) runOpts := spraycore.RunOptions{ Output: &buf, DefaultConfig: ".spray.yaml", @@ -145,13 +146,17 @@ func (c *Command) TestInjectProxy(args []string) []string { } func (c *Command) injectProxy(args []string) []string { - if c.Proxy == "" { + return c.injectProxyURL(args, c.Proxy) +} + +func (c *Command) injectProxyURL(args []string, proxy string) []string { + if proxy == "" { return args } if toolargs.HasFlag(args, "--proxy") { return args } - return append(args, "--proxy", c.Proxy) + return append(args, "--proxy", proxy) } func withDefaultNoBar(args []string) []string { diff --git a/tools/zombie/zombie.go b/tools/zombie/zombie.go index e1519375..8d55f5bc 100644 --- a/tools/zombie/zombie.go +++ b/tools/zombie/zombie.go @@ -4,14 +4,18 @@ import ( "bytes" "context" "fmt" + "net/url" "os" + "strings" aop "github.com/chainreactors/aiscan/aop" "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/tools/toolargs" + "github.com/chainreactors/proxyclient" sdkzombie "github.com/chainreactors/sdk/zombie" zombiecore "github.com/chainreactors/zombie/core" + zombiepkg "github.com/chainreactors/zombie/pkg" ) type Command struct { @@ -52,6 +56,11 @@ func (c *Command) Run(ctx context.Context, execution *commands.Execution) (_ any args := execution.Args args = c.resolveRelativePaths(args) args = ensureOutputDrain(args) + egress := commands.ResolveExecutionEgress(execution, c.Proxy) + proxyDial, err := proxyDialFor(egress.ProxyURL) + if err != nil { + return nil, err + } var buf bytes.Buffer if toolargs.BoolFlagEnabled(args, "--debug") { restoreDebug := telemetry.ActivateDebug(c.Logger) @@ -59,7 +68,8 @@ func (c *Command) Run(ctx context.Context, execution *commands.Execution) (_ any c.Logger.Debugf("zombie debug enabled") } runOpts := zombiecore.RunOptions{ - Output: &buf, + Output: &buf, + ProxyDial: proxyDial, } if err := zombiecore.RunWithArgs(ctx, args, runOpts); err != nil { if buf.Len() > 0 { @@ -71,6 +81,25 @@ func (c *Command) Run(ctx context.Context, execution *commands.Execution) (_ any return nil, nil } +func proxyDialFor(proxyURL string) (zombiepkg.DialFunc, error) { + proxyURL = strings.TrimSpace(proxyURL) + if proxyURL == "" { + return nil, nil + } + parsed, err := url.Parse(proxyURL) + if err != nil { + return nil, fmt.Errorf("zombie: invalid proxy %q: %w", proxyURL, err) + } + if parsed.Scheme == "" || parsed.Host == "" { + return nil, fmt.Errorf("zombie: invalid proxy %q: expected URL with scheme and host", proxyURL) + } + dial, err := proxyclient.NewClient(parsed) + if err != nil { + return nil, fmt.Errorf("zombie: initialize proxy %q: %w", proxyURL, err) + } + return dial.DialContext, nil +} + // zombie core only starts its result consumer when a file output is present, // while workers always publish to the result channel. Supply the system sink // for normal stdout-only runs so successful and failed attempts cannot deadlock. diff --git a/tools/zombie/zombie_test.go b/tools/zombie/zombie_test.go index 658ef0fc..be3fe7c0 100644 --- a/tools/zombie/zombie_test.go +++ b/tools/zombie/zombie_test.go @@ -26,6 +26,20 @@ func TestExecuteDebugActivatesTelemetryLogger(t *testing.T) { } } +func TestExecuteRejectsInvalidCallScopedProxy(t *testing.T) { + cmd := New(nil).WithProxy("http://startup.example:8080") + var output bytes.Buffer + _, err := cmd.Run(context.Background(), &commands.Execution{ + Args: []string{"--help"}, + Env: []string{"ALL_PROXY=not-a-proxy"}, + Stdout: &output, + Stderr: &output, + }) + if err == nil || !strings.Contains(err.Error(), "invalid proxy") { + t.Fatalf("Run() error = %v, want invalid call-scoped proxy error", err) + } +} + func TestResolveRelativePathsOnlyRewritesZombieFileFlags(t *testing.T) { dir := t.TempDir() cmd := New(nil) diff --git a/web/frontend/src/compat/ioa.tsx b/web/frontend/src/compat/ioa.tsx deleted file mode 100644 index a2ad8936..00000000 --- a/web/frontend/src/compat/ioa.tsx +++ /dev/null @@ -1,70 +0,0 @@ -import { Badge } from '@cyber/ui' -import { cn } from '@cyber/theme' -import { - GraphPanel as CyberGraphPanel, - MessageContent as CyberMessageContent, - type GraphPanelProps as CyberGraphPanelProps, - type MessageContentProps as CyberMessageContentProps, -} from '../../cyber-ui/packages/ioa/src' - -export * from '../../cyber-ui/packages/ioa/src' - -export interface GraphPanelProps extends CyberGraphPanelProps { - title?: string -} - -export function GraphPanel({ title: _title, ...props }: GraphPanelProps) { - return -} - -export interface MessageContentProps extends CyberMessageContentProps { - showFrontMatter?: boolean - showType?: boolean -} - -export function MessageContent({ - showFrontMatter: _showFrontMatter, - showType: _showType, - ...props -}: MessageContentProps) { - return -} - -export interface MessageFrontMatterProps { - content: unknown - meta?: Record - className?: string -} - -export function MessageFrontMatter({ content, meta, className }: MessageFrontMatterProps) { - const record = content && typeof content === 'object' && !Array.isArray(content) - ? content as Record - : null - const contentType = typeof record?.type === 'string' ? record.type : '' - const metaKind = typeof meta?.kind === 'string' ? meta.kind : '' - const metaLabels = Array.isArray(meta?.labels) - ? meta.labels.filter((label): label is string => typeof label === 'string') - : [] - - if (!contentType && !metaKind && metaLabels.length === 0) return null - - return ( -
- {contentType && ( - - {contentType} - - )} - {metaKind && ( - - {metaKind} - - )} - {metaLabels.map(label => ( - - {label} - - ))} -
- ) -} diff --git a/web/frontend/src/components/IOAConsole.tsx b/web/frontend/src/components/IOAConsole.tsx index b0b31bb6..c70f98e3 100644 --- a/web/frontend/src/components/IOAConsole.tsx +++ b/web/frontend/src/components/IOAConsole.tsx @@ -5,7 +5,6 @@ import { GraphPanel, HandoffCard, MessageContent, - MessageFrontMatter, detectContentType, messageTitle, type ForumThread, @@ -204,7 +203,6 @@ export default function IOAConsole({ open, onClose, initialSpaceID, initialMessa selectedMessageId={selectedMessageID} onSelectMessage={setSelectedMessageID} mode="dialog" - title="IOA" /> ) : (
@@ -328,8 +326,6 @@ function MessageInspector({ )}
@@ -338,6 +334,47 @@ function MessageInspector({ ) } +function MessageFrontMatter({ + content, + meta, + className, +}: { + content: unknown + meta?: Record + className?: string +}) { + const record = content && typeof content === 'object' && !Array.isArray(content) + ? content as Record + : null + const contentType = typeof record?.type === 'string' ? record.type : '' + const metaKind = typeof meta?.kind === 'string' ? meta.kind : '' + const metaLabels = Array.isArray(meta?.labels) + ? meta.labels.filter((label): label is string => typeof label === 'string') + : [] + + if (!contentType && !metaKind && metaLabels.length === 0) return null + + return ( +
+ {contentType && ( + + {contentType} + + )} + {metaKind && ( + + {metaKind} + + )} + {metaLabels.map(label => ( + + {label} + + ))} +
+ ) +} + function formatMessageTime(value: string) { const timestamp = Date.parse(value) if (!Number.isFinite(timestamp)) return value || '—' diff --git a/web/frontend/src/gen/types/chat_pb.ts b/web/frontend/src/gen/types/chat_pb.ts index ea9534c5..e09ddd5b 100644 --- a/web/frontend/src/gen/types/chat_pb.ts +++ b/web/frontend/src/gen/types/chat_pb.ts @@ -2,8 +2,8 @@ // @generated from file types/chat.proto (package aiscan.chat, syntax proto3) /* eslint-disable */ -import type { GenFile, GenMessage } from "@bufbuild/protobuf/codegenv2"; -import { fileDesc, messageDesc } from "@bufbuild/protobuf/codegenv2"; +import type { GenEnum, GenFile, GenMessage } from "@bufbuild/protobuf/codegenv2"; +import { enumDesc, fileDesc, messageDesc } from "@bufbuild/protobuf/codegenv2"; import type { Rejection, Session } from "../../../cyber-ui/packages/aop/src/gen/aop/chat_pb.js"; import { file_aop_chat } from "../../../cyber-ui/packages/aop/src/gen/aop/chat_pb.js"; import type { CommandSpec } from "./command_pb.js"; @@ -16,7 +16,53 @@ import type { Message } from "@bufbuild/protobuf"; * Describes the file types/chat.proto. */ export const file_types_chat: GenFile = /*@__PURE__*/ - fileDesc("ChB0eXBlcy9jaGF0LnByb3RvEgthaXNjYW4uY2hhdCK0AQoNU2Vzc2lvblJlY29yZBIdCgdzZXNzaW9uGAEgASgLMgwuYW9wLlNlc3Npb24SEgoKYWdlbnRfbmFtZRgCIAEoCRIQCghzY2FuX2lkcxgDIAMoCRIuCgpjcmVhdGVkX2F0GAQgASgLMhouZ29vZ2xlLnByb3RvYnVmLlRpbWVzdGFtcBIuCgp1cGRhdGVkX2F0GAUgASgLMhouZ29vZ2xlLnByb3RvYnVmLlRpbWVzdGFtcCJSChNMaXN0U2Vzc2lvbnNSZXF1ZXN0EhQKDGFmdGVyX2N1cnNvchgBIAEoCRINCgVsaW1pdBgCIAEoDRIWCg5pbmNsdWRlX2Nsb3NlZBgDIAEoCCJZChRMaXN0U2Vzc2lvbnNSZXNwb25zZRIsCghzZXNzaW9ucxgBIAMoCzIaLmFpc2Nhbi5jaGF0LlNlc3Npb25SZWNvcmQSEwoLbmV4dF9jdXJzb3IYAiABKAkiJwoRR2V0U2Vzc2lvblJlcXVlc3QSEgoKc2Vzc2lvbl9pZBgBIAEoCSJBChJHZXRTZXNzaW9uUmVzcG9uc2USKwoHc2Vzc2lvbhgBIAEoCzIaLmFpc2Nhbi5jaGF0LlNlc3Npb25SZWNvcmQiZAoTUmVzZXRTZXNzaW9uUmVxdWVzdBISCgpyZXF1ZXN0X2lkGAEgASgJEhIKCnNlc3Npb25faWQYAiABKAkSFgoObmV3X3Nlc3Npb25faWQYAyABKAkSDQoFdGl0bGUYBCABKAkiYgoTUmVzZXRTZXNzaW9uUmVjZWlwdBIeCghwcmV2aW91cxgBIAEoCzIMLmFvcC5TZXNzaW9uEisKB2N1cnJlbnQYAiABKAsyGi5haXNjYW4uY2hhdC5TZXNzaW9uUmVjb3JkIo8BChRSZXNldFNlc3Npb25SZXNwb25zZRISCgpyZXF1ZXN0X2lkGAEgASgJEjQKCGFjY2VwdGVkGAIgASgLMiAuYWlzY2FuLmNoYXQuUmVzZXRTZXNzaW9uUmVjZWlwdEgAEiIKCHJlamVjdGVkGAMgASgLMg4uYW9wLlJlamVjdGlvbkgAQgkKB291dGNvbWUiPgoURGVsZXRlU2Vzc2lvblJlcXVlc3QSEgoKcmVxdWVzdF9pZBgBIAEoCRISCgpzZXNzaW9uX2lkGAIgASgJInwKFURlbGV0ZVNlc3Npb25SZXNwb25zZRISCgpyZXF1ZXN0X2lkGAEgASgJEiAKCGFjY2VwdGVkGAIgASgLMgwuYW9wLlNlc3Npb25IABIiCghyZWplY3RlZBgDIAEoCzIOLmFvcC5SZWplY3Rpb25IAEIJCgdvdXRjb21lIikKE0xpc3RDb21tYW5kc1JlcXVlc3QSEgoKc2Vzc2lvbl9pZBgBIAEoCSJFChRMaXN0Q29tbWFuZHNSZXNwb25zZRItCghjb21tYW5kcxgBIAMoCzIbLmFpc2Nhbi5jb21tYW5kLkNvbW1hbmRTcGVjQjFaL2dpdGh1Yi5jb20vY2hhaW5yZWFjdG9ycy9haXNjYW4vcGtnL3R5cGVzO3R5cGVzYgZwcm90bzM", [file_aop_chat, file_types_command, file_google_protobuf_timestamp]); + fileDesc("ChB0eXBlcy9jaGF0LnByb3RvEgthaXNjYW4uY2hhdCKDAQoOU2Vzc2lvbkhpc3RvcnkSLgoEbW9kZRgBIAEoDjIgLmFpc2Nhbi5jaGF0LlNlc3Npb25IaXN0b3J5Lk1vZGUiQQoETW9kZRIUChBNT0RFX1VOU1BFQ0lGSUVEEAASEAoMTU9ERV9JTkhFUklUEAESEQoNTU9ERV9TTkFQU0hPVBACIrQBCg1TZXNzaW9uUmVjb3JkEh0KB3Nlc3Npb24YASABKAsyDC5hb3AuU2Vzc2lvbhISCgphZ2VudF9uYW1lGAIgASgJEhAKCHNjYW5faWRzGAMgAygJEi4KCmNyZWF0ZWRfYXQYBCABKAsyGi5nb29nbGUucHJvdG9idWYuVGltZXN0YW1wEi4KCnVwZGF0ZWRfYXQYBSABKAsyGi5nb29nbGUucHJvdG9idWYuVGltZXN0YW1wIlIKE0xpc3RTZXNzaW9uc1JlcXVlc3QSFAoMYWZ0ZXJfY3Vyc29yGAEgASgJEg0KBWxpbWl0GAIgASgNEhYKDmluY2x1ZGVfY2xvc2VkGAMgASgIIlkKFExpc3RTZXNzaW9uc1Jlc3BvbnNlEiwKCHNlc3Npb25zGAEgAygLMhouYWlzY2FuLmNoYXQuU2Vzc2lvblJlY29yZBITCgtuZXh0X2N1cnNvchgCIAEoCSInChFHZXRTZXNzaW9uUmVxdWVzdBISCgpzZXNzaW9uX2lkGAEgASgJIkEKEkdldFNlc3Npb25SZXNwb25zZRIrCgdzZXNzaW9uGAEgASgLMhouYWlzY2FuLmNoYXQuU2Vzc2lvblJlY29yZCJkChNSZXNldFNlc3Npb25SZXF1ZXN0EhIKCnJlcXVlc3RfaWQYASABKAkSEgoKc2Vzc2lvbl9pZBgCIAEoCRIWCg5uZXdfc2Vzc2lvbl9pZBgDIAEoCRINCgV0aXRsZRgEIAEoCSJiChNSZXNldFNlc3Npb25SZWNlaXB0Eh4KCHByZXZpb3VzGAEgASgLMgwuYW9wLlNlc3Npb24SKwoHY3VycmVudBgCIAEoCzIaLmFpc2Nhbi5jaGF0LlNlc3Npb25SZWNvcmQijwEKFFJlc2V0U2Vzc2lvblJlc3BvbnNlEhIKCnJlcXVlc3RfaWQYASABKAkSNAoIYWNjZXB0ZWQYAiABKAsyIC5haXNjYW4uY2hhdC5SZXNldFNlc3Npb25SZWNlaXB0SAASIgoIcmVqZWN0ZWQYAyABKAsyDi5hb3AuUmVqZWN0aW9uSABCCQoHb3V0Y29tZSI+ChREZWxldGVTZXNzaW9uUmVxdWVzdBISCgpyZXF1ZXN0X2lkGAEgASgJEhIKCnNlc3Npb25faWQYAiABKAkifAoVRGVsZXRlU2Vzc2lvblJlc3BvbnNlEhIKCnJlcXVlc3RfaWQYASABKAkSIAoIYWNjZXB0ZWQYAiABKAsyDC5hb3AuU2Vzc2lvbkgAEiIKCHJlamVjdGVkGAMgASgLMg4uYW9wLlJlamVjdGlvbkgAQgkKB291dGNvbWUiKQoTTGlzdENvbW1hbmRzUmVxdWVzdBISCgpzZXNzaW9uX2lkGAEgASgJIkUKFExpc3RDb21tYW5kc1Jlc3BvbnNlEi0KCGNvbW1hbmRzGAEgAygLMhsuYWlzY2FuLmNvbW1hbmQuQ29tbWFuZFNwZWNCMVovZ2l0aHViLmNvbS9jaGFpbnJlYWN0b3JzL2Fpc2Nhbi9wa2cvdHlwZXM7dHlwZXNiBnByb3RvMw", [file_aop_chat, file_types_command, file_google_protobuf_timestamp]); + +/** + * SessionHistory is persisted as an AOP event extension. It makes transcript + * inheritance explicit without changing the shared AOP protocol schema. + * + * @generated from message aiscan.chat.SessionHistory + */ +export type SessionHistory = Message<"aiscan.chat.SessionHistory"> & { + /** + * @generated from field: aiscan.chat.SessionHistory.Mode mode = 1; + */ + mode: SessionHistory_Mode; +}; + +/** + * Describes the message aiscan.chat.SessionHistory. + * Use `create(SessionHistorySchema)` to create a new message. + */ +export const SessionHistorySchema: GenMessage = /*@__PURE__*/ + messageDesc(file_types_chat, 0); + +/** + * @generated from enum aiscan.chat.SessionHistory.Mode + */ +export enum SessionHistory_Mode { + /** + * @generated from enum value: MODE_UNSPECIFIED = 0; + */ + UNSPECIFIED = 0, + + /** + * @generated from enum value: MODE_INHERIT = 1; + */ + INHERIT = 1, + + /** + * @generated from enum value: MODE_SNAPSHOT = 2; + */ + SNAPSHOT = 2, +} + +/** + * Describes the enum aiscan.chat.SessionHistory.Mode. + */ +export const SessionHistory_ModeSchema: GenEnum = /*@__PURE__*/ + enumDesc(file_types_chat, 0, 0); /** * @generated from message aiscan.chat.SessionRecord @@ -53,7 +99,7 @@ export type SessionRecord = Message<"aiscan.chat.SessionRecord"> & { * Use `create(SessionRecordSchema)` to create a new message. */ export const SessionRecordSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_types_chat, 0); + messageDesc(file_types_chat, 1); /** * @generated from message aiscan.chat.ListSessionsRequest @@ -80,7 +126,7 @@ export type ListSessionsRequest = Message<"aiscan.chat.ListSessionsRequest"> & { * Use `create(ListSessionsRequestSchema)` to create a new message. */ export const ListSessionsRequestSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_types_chat, 1); + messageDesc(file_types_chat, 2); /** * @generated from message aiscan.chat.ListSessionsResponse @@ -102,7 +148,7 @@ export type ListSessionsResponse = Message<"aiscan.chat.ListSessionsResponse"> & * Use `create(ListSessionsResponseSchema)` to create a new message. */ export const ListSessionsResponseSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_types_chat, 2); + messageDesc(file_types_chat, 3); /** * @generated from message aiscan.chat.GetSessionRequest @@ -119,7 +165,7 @@ export type GetSessionRequest = Message<"aiscan.chat.GetSessionRequest"> & { * Use `create(GetSessionRequestSchema)` to create a new message. */ export const GetSessionRequestSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_types_chat, 3); + messageDesc(file_types_chat, 4); /** * @generated from message aiscan.chat.GetSessionResponse @@ -136,7 +182,7 @@ export type GetSessionResponse = Message<"aiscan.chat.GetSessionResponse"> & { * Use `create(GetSessionResponseSchema)` to create a new message. */ export const GetSessionResponseSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_types_chat, 4); + messageDesc(file_types_chat, 5); /** * @generated from message aiscan.chat.ResetSessionRequest @@ -168,7 +214,7 @@ export type ResetSessionRequest = Message<"aiscan.chat.ResetSessionRequest"> & { * Use `create(ResetSessionRequestSchema)` to create a new message. */ export const ResetSessionRequestSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_types_chat, 5); + messageDesc(file_types_chat, 6); /** * @generated from message aiscan.chat.ResetSessionReceipt @@ -190,7 +236,7 @@ export type ResetSessionReceipt = Message<"aiscan.chat.ResetSessionReceipt"> & { * Use `create(ResetSessionReceiptSchema)` to create a new message. */ export const ResetSessionReceiptSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_types_chat, 6); + messageDesc(file_types_chat, 7); /** * @generated from message aiscan.chat.ResetSessionResponse @@ -224,7 +270,7 @@ export type ResetSessionResponse = Message<"aiscan.chat.ResetSessionResponse"> & * Use `create(ResetSessionResponseSchema)` to create a new message. */ export const ResetSessionResponseSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_types_chat, 7); + messageDesc(file_types_chat, 8); /** * @generated from message aiscan.chat.DeleteSessionRequest @@ -246,7 +292,7 @@ export type DeleteSessionRequest = Message<"aiscan.chat.DeleteSessionRequest"> & * Use `create(DeleteSessionRequestSchema)` to create a new message. */ export const DeleteSessionRequestSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_types_chat, 8); + messageDesc(file_types_chat, 9); /** * @generated from message aiscan.chat.DeleteSessionResponse @@ -280,7 +326,7 @@ export type DeleteSessionResponse = Message<"aiscan.chat.DeleteSessionResponse"> * Use `create(DeleteSessionResponseSchema)` to create a new message. */ export const DeleteSessionResponseSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_types_chat, 9); + messageDesc(file_types_chat, 10); /** * @generated from message aiscan.chat.ListCommandsRequest @@ -297,7 +343,7 @@ export type ListCommandsRequest = Message<"aiscan.chat.ListCommandsRequest"> & { * Use `create(ListCommandsRequestSchema)` to create a new message. */ export const ListCommandsRequestSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_types_chat, 10); + messageDesc(file_types_chat, 11); /** * @generated from message aiscan.chat.ListCommandsResponse @@ -314,4 +360,4 @@ export type ListCommandsResponse = Message<"aiscan.chat.ListCommandsResponse"> & * Use `create(ListCommandsResponseSchema)` to create a new message. */ export const ListCommandsResponseSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_types_chat, 11); + messageDesc(file_types_chat, 12); diff --git a/web/frontend/tsconfig.json b/web/frontend/tsconfig.json index a07c9661..18b5f3ee 100644 --- a/web/frontend/tsconfig.json +++ b/web/frontend/tsconfig.json @@ -27,7 +27,7 @@ "@cyber/cstx": ["./cyber-ui/packages/cstx/src"], "@cyber/cstx-easm": ["./cyber-ui/packages/cstx-easm/src"], "@cyber/viewer": ["./cyber-ui/packages/viewer/src"], - "@cyber/ioa": ["./src/compat/ioa.tsx"] + "@cyber/ioa": ["./cyber-ui/packages/ioa/src"] } }, "include": [ diff --git a/web/frontend/vite.config.ts b/web/frontend/vite.config.ts index b5638b1e..71dde1fe 100644 --- a/web/frontend/vite.config.ts +++ b/web/frontend/vite.config.ts @@ -4,7 +4,7 @@ import path from 'path' const backendURL = process.env.AISCAN_BACKEND_URL || 'http://127.0.0.1:8080' -// Design-system primitives + the terminal view are consumed from the cyber-ui +// Shared UI and IOA components are consumed directly from the cyber-ui // submodule (single source of truth for what aiscan contributes upstream). The // remaining composite views (markdown/viewer) stay vendored under @/ because // aiscan still diverges them. @@ -23,7 +23,7 @@ export default defineConfig({ '@cyber/cstx': path.resolve(cyberUI, 'cstx/src'), '@cyber/cstx-easm': path.resolve(cyberUI, 'cstx-easm/src'), '@cyber/viewer': path.resolve(cyberUI, 'viewer/src'), - '@cyber/ioa': path.resolve(__dirname, './src/compat/ioa.tsx'), + '@cyber/ioa': path.resolve(cyberUI, 'ioa/src'), }, }, server: {