From bdda0af411dad8434e7b45b467d3721430fca1dc Mon Sep 17 00:00:00 2001 From: M09Ic Date: Mon, 3 Aug 2026 12:59:56 +0800 Subject: [PATCH] fix(session): unify JSONL persistence and rotate continuations --- agent/agent.go | 3 + agent/aop_emit.go | 34 +- agent/checkpoint.go | 231 ---------- agent/checkpoint_test.go | 170 -------- agent/loop.go | 8 +- agent/subagent.go | 3 +- agent/subagent_test.go | 5 +- agent/types.go | 38 +- aop/helpers.go | 4 + aop/tool/artifact.go | 10 + aop/tool/protocol.pb.go | 24 +- cmd/aiscan/cli.go | 54 ++- cmd/aiscan/cli_test.go | 54 +++ cmd/aiscan/setup.go | 9 +- cmd/aiscan/web_full.go | 15 +- cmd/aiscan/web_full_test.go | 28 +- core/config/options.go | 10 +- core/output/artifact_stream.go | 83 ---- core/output/artifact_stream_test.go | 49 --- core/output/jsonl.go | 188 ++++++++ core/output/jsonl_test.go | 83 ++++ core/output/{timeline.go => render.go} | 167 +++----- core/output/render_test.go | 136 ++++++ core/output/timeline_test.go | 65 --- core/output/tool_data.go | 38 -- core/output/types.go | 2 +- core/tool/context.go | 6 +- docs/mechanisms.md | 6 +- docs/protocol-architecture.md | 4 +- docs/scan.md | 4 +- examples/acp/server/main.go | 15 +- pkg/commands/factory.go | 5 +- pkg/node/agent.go | 3 +- pkg/node/connection.go | 5 +- pkg/node/proto_connection.go | 21 +- pkg/node/toolnode.go | 101 ++--- pkg/node/toolnode_test.go | 28 +- pkg/runner/app.go | 97 ++++- pkg/runner/app_test.go | 87 ++++ pkg/runner/application_builder.go | 7 +- pkg/runner/application_config.go | 1 + pkg/runner/local_repl.go | 2 +- pkg/runner/provider_config_test.go | 3 +- pkg/runner/remote_repl.go | 2 +- pkg/runner/runner.go | 210 ++++++--- pkg/runner/runner_test.go | 401 ++++++++++++++++++ pkg/runner/runtime_session.go | 370 ++++++++++++++-- pkg/runner/runtime_session_test.go | 68 ++- pkg/runner/session_jsonl.go | 139 ++++++ pkg/runner/session_jsonl_test.go | 85 ++++ pkg/runner/tool_call.go | 36 +- pkg/runner/tool_call_test.go | 15 +- pkg/tui/commands.go | 15 + pkg/tui/console.go | 56 ++- pkg/tui/console_test.go | 162 +++++-- pkg/tui/controller.go | 9 + pkg/tui/remote_console.go | 22 +- pkg/tui/remote_console_test.go | 45 +- pkg/web/service/agents_mux.go | 41 +- pkg/web/service/agents_test.go | 14 +- pkg/web/service/artifacts.go | 4 +- pkg/web/service/artifacts_native.go | 9 +- pkg/web/service/artifacts_test.go | 8 +- pkg/web/service/scan.go | 3 +- tools/gogo/gogo.go | 10 +- tools/gogo/register.go | 2 +- tools/katana/katana.go | 10 +- tools/katana/register.go | 2 +- tools/neutron/neutron.go | 10 +- tools/neutron/register.go | 2 +- tools/proton/command.go | 10 +- tools/proton/register.go | 2 +- tools/register_command.go | 4 +- .../register_command_full_integration_test.go | 8 +- tools/register_command_full_test.go | 9 +- tools/register_command_integration_test.go | 12 +- tools/register_command_test.go | 77 +++- tools/scan/command.go | 56 +-- tools/scan/command_test.go | 77 ++-- tools/scan/options.go | 7 +- tools/spray/register.go | 2 +- tools/spray/spray.go | 10 +- tools/toolargs/base.go | 47 +- tools/zombie/register.go | 2 +- tools/zombie/zombie.go | 7 +- web/frontend/cyber-ui | 2 +- 86 files changed, 2647 insertions(+), 1341 deletions(-) delete mode 100644 agent/checkpoint.go delete mode 100644 agent/checkpoint_test.go create mode 100644 aop/tool/artifact.go delete mode 100644 core/output/artifact_stream.go delete mode 100644 core/output/artifact_stream_test.go create mode 100644 core/output/jsonl.go create mode 100644 core/output/jsonl_test.go rename core/output/{timeline.go => render.go} (68%) create mode 100644 core/output/render_test.go delete mode 100644 core/output/timeline_test.go delete mode 100644 core/output/tool_data.go create mode 100644 pkg/runner/runner_test.go create mode 100644 pkg/runner/session_jsonl.go create mode 100644 pkg/runner/session_jsonl_test.go diff --git a/agent/agent.go b/agent/agent.go index 2e4ffeff..a616869a 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -281,6 +281,9 @@ func (a *Agent) LoadMessages(messages []*aop.Message) { a.mu.Lock() defer a.mu.Unlock() a.state.Messages = append([]*aop.Message(nil), messages...) + if a.Cfg.emitter != nil { + a.Cfg.emitter.observeMessages(messages) + } } func (a *Agent) validateContinue() error { diff --git a/agent/aop_emit.go b/agent/aop_emit.go index 12544cf0..6164f5d6 100644 --- a/agent/aop_emit.go +++ b/agent/aop_emit.go @@ -2,6 +2,8 @@ package agent import ( "fmt" + "strconv" + "strings" "sync/atomic" aop "github.com/chainreactors/aiscan/aop" @@ -18,15 +20,8 @@ const ( statusLLMRequest = "llm_request" ) -// EventEmitter is the narrow event sink an agent needs — callers may wrap a -// bus with stamping/routing middleware (e.g. the runner's sessionEmitter) -// instead of handing over a raw *eventbus.Bus. -type EventEmitter interface { - Emit(*aop.Event) -} - type aopEmitter struct { - bus EventEmitter + bus aop.EventEmitter agentName string sessionID string turnID string @@ -41,7 +36,7 @@ type emitState struct { messageSeq atomic.Int64 } -func newAOPEmitter(bus EventEmitter, agentName, sessionID, parentSessionID, parentToolCallID string, detail *types.DelegationDetail, msgCounter int64) *aopEmitter { +func newAOPEmitter(bus aop.EventEmitter, agentName, sessionID, parentSessionID, parentToolCallID string, detail *types.DelegationDetail, msgCounter int64) *aopEmitter { em := &aopEmitter{ bus: bus, agentName: agentName, sessionID: sessionID, parentSessionID: parentSessionID, parentToolCallID: parentToolCallID, @@ -82,6 +77,27 @@ func (e *aopEmitter) allocMessageID() string { func (e *aopEmitter) messageCounter() int64 { return e.state.messageSeq.Load() } +func (e *aopEmitter) observeMessages(messages []*aop.Message) { + if e == nil || e.state == nil { + return + } + var observed int64 + for _, message := range messages { + if message == nil || !strings.HasPrefix(message.Id, "m-") { + continue + } + sequence, err := strconv.ParseInt(strings.TrimPrefix(message.Id, "m-"), 10, 64) + if err == nil && sequence > observed { + observed = sequence + } + } + for current := e.state.messageSeq.Load(); observed > current; current = e.state.messageSeq.Load() { + if e.state.messageSeq.CompareAndSwap(current, observed) { + return + } + } +} + func (e *aopEmitter) sessionStart(model string) { event := &aop.Event{Payload: &aop.Event_SessionStarted{SessionStarted: &aop.SessionStarted{ Model: model, ParentSessionId: e.parentSessionID, ParentToolCallId: e.parentToolCallID, diff --git a/agent/checkpoint.go b/agent/checkpoint.go deleted file mode 100644 index fb8f77ab..00000000 --- a/agent/checkpoint.go +++ /dev/null @@ -1,231 +0,0 @@ -package agent - -import ( - "encoding/json" - "fmt" - "os" - "path/filepath" - "sort" - "time" - - aop "github.com/chainreactors/aiscan/aop" - "google.golang.org/protobuf/encoding/protojson" - "google.golang.org/protobuf/proto" -) - -type CheckpointData struct { - Version int `json:"version"` - CreatedAt time.Time `json:"created_at"` - UpdatedAt time.Time `json:"updated_at"` - Model string `json:"model,omitempty"` - Provider string `json:"provider,omitempty"` - Messages []*aop.Message `json:"messages"` - // MessageCounter resumes AOP message_id allocation ("m-") after restore. - MessageCounter int64 `json:"message_counter,omitempty"` -} - -type CheckpointInfo struct { - Path string - CreatedAt time.Time - UpdatedAt time.Time - ModTime time.Time - Model string - Provider string - Messages int -} - -const checkpointVersion = 1 - -func SaveCheckpoint(dir string, data *CheckpointData) error { - if err := os.MkdirAll(dir, 0o755); err != nil { - return fmt.Errorf("create session dir: %w", err) - } - now := time.Now() - if data.CreatedAt.IsZero() { - data.CreatedAt = now - } - data.UpdatedAt = now - data.Version = checkpointVersion - data.Messages = sanitizeMessagesForSave(data.Messages) - - raw, err := json.MarshalIndent(checkpointJSON{ - Version: data.Version, - CreatedAt: data.CreatedAt, - UpdatedAt: data.UpdatedAt, - Model: data.Model, - Provider: data.Provider, - Messages: marshalMessages(data.Messages), - MessageCounter: data.MessageCounter, - }, "", " ") - if err != nil { - return fmt.Errorf("marshal session: %w", err) - } - - ts := now.Format("20060102-150405") - tsPath := filepath.Join(dir, fmt.Sprintf("session-%s.json", ts)) - if err := os.WriteFile(tsPath, raw, 0o644); err != nil { - return fmt.Errorf("write session file: %w", err) - } - - return nil -} - -func LoadCheckpoint(path string) (*CheckpointData, error) { - raw, err := os.ReadFile(path) - if err != nil { - return nil, fmt.Errorf("read session file: %w", err) - } - var data checkpointJSON - if err := json.Unmarshal(raw, &data); err != nil { - return nil, fmt.Errorf("parse session file: %w", err) - } - messages, err := unmarshalMessages(data.Messages) - if err != nil { - return nil, fmt.Errorf("parse session messages: %w", err) - } - return &CheckpointData{ - Version: data.Version, - CreatedAt: data.CreatedAt, - UpdatedAt: data.UpdatedAt, - Model: data.Model, - Provider: data.Provider, - Messages: messages, - MessageCounter: data.MessageCounter, - }, nil -} - -// checkpointJSON is the on-disk envelope. Messages are stored as proto-JSON so -// the file format mirrors the AOP truth instead of a vendor shape. -type checkpointJSON struct { - Version int `json:"version"` - CreatedAt time.Time `json:"created_at"` - UpdatedAt time.Time `json:"updated_at"` - Model string `json:"model,omitempty"` - Provider string `json:"provider,omitempty"` - Messages []json.RawMessage `json:"messages"` - MessageCounter int64 `json:"message_counter,omitempty"` -} - -func marshalMessages(messages []*aop.Message) []json.RawMessage { - out := make([]json.RawMessage, 0, len(messages)) - for _, m := range messages { - raw, err := protojson.Marshal(m) - if err != nil { - continue - } - out = append(out, raw) - } - return out -} - -func unmarshalMessages(raw []json.RawMessage) ([]*aop.Message, error) { - out := make([]*aop.Message, 0, len(raw)) - for _, data := range raw { - msg := new(aop.Message) - if err := protojson.Unmarshal(data, msg); err != nil { - return nil, err - } - out = append(out, msg) - } - return out, nil -} - -type checkpointMeta struct { - CreatedAt time.Time `json:"created_at"` - UpdatedAt time.Time `json:"updated_at"` - Model string `json:"model,omitempty"` - Provider string `json:"provider,omitempty"` - Messages []json.RawMessage `json:"messages"` -} - -func ListCheckpoints(dir string) ([]CheckpointInfo, error) { - entries, err := os.ReadDir(dir) - if err != nil { - if os.IsNotExist(err) { - return nil, nil - } - return nil, fmt.Errorf("read session dir: %w", err) - } - sessions := make([]CheckpointInfo, 0, len(entries)) - for _, entry := range entries { - if entry.IsDir() { - continue - } - if matched, _ := filepath.Match("session-*.json", entry.Name()); !matched { - continue - } - path := filepath.Join(dir, entry.Name()) - raw, err := os.ReadFile(path) - if err != nil { - continue - } - var meta checkpointMeta - if err := json.Unmarshal(raw, &meta); err != nil { - continue - } - fi, _ := entry.Info() - modTime := time.Time{} - if fi != nil { - modTime = fi.ModTime() - } - sessions = append(sessions, CheckpointInfo{ - Path: path, - CreatedAt: meta.CreatedAt, - UpdatedAt: meta.UpdatedAt, - ModTime: modTime, - Model: meta.Model, - Provider: meta.Provider, - Messages: len(meta.Messages), - }) - } - sort.Slice(sessions, func(i, j int) bool { - left := sessions[i].SortTime() - right := sessions[j].SortTime() - if left.Equal(right) { - return sessions[i].Path > sessions[j].Path - } - return left.After(right) - }) - return sessions, nil -} - -func (s CheckpointInfo) SortTime() time.Time { - switch { - case !s.UpdatedAt.IsZero(): - return s.UpdatedAt - case !s.CreatedAt.IsZero(): - return s.CreatedAt - default: - return s.ModTime - } -} - -// sanitizeMessagesForSave strips binary media parts before persisting: an -// image is re-fetchable context, not history worth 20 MiB of JSON. Text and -// tool call/result structure is preserved. -func sanitizeMessagesForSave(messages []*aop.Message) []*aop.Message { - out := make([]*aop.Message, len(messages)) - for i, m := range messages { - hasMedia := false - for _, part := range m.Content { - if part.GetMedia() != nil { - hasMedia = true - break - } - } - if !hasMedia { - out[i] = m - continue - } - filtered := make([]*aop.Content, 0, len(m.Content)) - for _, part := range m.Content { - if part.GetMedia() == nil { - filtered = append(filtered, part) - } - } - cp := proto.CloneOf(m) - cp.Content = filtered - out[i] = cp - } - return out -} diff --git a/agent/checkpoint_test.go b/agent/checkpoint_test.go deleted file mode 100644 index 6ecccad0..00000000 --- a/agent/checkpoint_test.go +++ /dev/null @@ -1,170 +0,0 @@ -package agent - -import ( - "encoding/json" - "os" - "path/filepath" - "testing" - "time" - - "github.com/chainreactors/aiscan/agent/provider" - aop "github.com/chainreactors/aiscan/aop" -) - -func TestSaveAndLoadCheckpoint(t *testing.T) { - dir := t.TempDir() - - content := "hello world" - toolArgs := `{"cmd":"ls"}` - messages := []*aop.Message{ - textMessage("user", content), - { - Role: "assistant", - Content: []*aop.Content{ - aop.Text(content), - toolCallContent("tc1", "bash", toolArgs), - }, - }, - toolResultMessage("tc1", content), - } - - data := &CheckpointData{ - Model: "gpt-4o", - Provider: "openai", - Messages: messages, - } - if err := SaveCheckpoint(dir, data); err != nil { - t.Fatalf("SaveCheckpoint: %v", err) - } - - if _, err := os.Stat(filepath.Join(dir, "latest.json")); !os.IsNotExist(err) { - t.Fatalf("latest.json should not be written, err=%v", err) - } - - sessions, err := ListCheckpoints(dir) - if err != nil { - t.Fatalf("ListCheckpoints: %v", err) - } - if len(sessions) != 1 { - t.Fatalf("sessions len = %d, want 1", len(sessions)) - } - - loaded, err := LoadCheckpoint(sessions[0].Path) - if err != nil { - t.Fatalf("LoadCheckpoint: %v", err) - } - if loaded.Version != checkpointVersion { - t.Errorf("version = %d, want %d", loaded.Version, checkpointVersion) - } - if loaded.Model != "gpt-4o" { - t.Errorf("model = %q, want %q", loaded.Model, "gpt-4o") - } - if len(loaded.Messages) != 3 { - t.Fatalf("messages len = %d, want 3", len(loaded.Messages)) - } - if loaded.Messages[0].Role != "user" || provider.MessageText(loaded.Messages[0]) != "hello world" { - t.Errorf("message[0] = %+v", loaded.Messages[0]) - } - if calls := provider.MessageToolCalls(loaded.Messages[1]); len(calls) != 1 || calls[0].Name != "bash" { - t.Errorf("message[1] tool_calls = %+v", calls) - } - if r := provider.MessageToolResult(loaded.Messages[2]); r == nil || r.CallId != "tc1" { - t.Errorf("message[2] tool result = %+v, want call id tc1", r) - } - - entries, _ := os.ReadDir(dir) - found := false - for _, e := range entries { - if matched, _ := filepath.Match("session-*.json", e.Name()); matched { - found = true - } - } - if !found { - t.Error("timestamped session file not found") - } -} - -func TestListCheckpointsSortsNewestFirst(t *testing.T) { - dir := t.TempDir() - oldTime := time.Date(2026, 7, 12, 10, 0, 0, 0, time.UTC) - newTime := time.Date(2026, 7, 13, 10, 0, 0, 0, time.UTC) - writeSessionFile(t, filepath.Join(dir, "session-old.json"), CheckpointData{ - Version: checkpointVersion, - UpdatedAt: oldTime, - Model: "old", - Messages: []*aop.Message{textMessage("user", "old")}, - }) - writeSessionFile(t, filepath.Join(dir, "session-new.json"), CheckpointData{ - Version: checkpointVersion, - UpdatedAt: newTime, - Model: "new", - Messages: []*aop.Message{textMessage("user", "new")}, - }) - writeSessionFile(t, filepath.Join(dir, "latest.json"), CheckpointData{ - Version: checkpointVersion, - UpdatedAt: newTime.Add(time.Hour), - Model: "ignored", - Messages: []*aop.Message{textMessage("user", "ignored")}, - }) - - sessions, err := ListCheckpoints(dir) - if err != nil { - t.Fatalf("ListCheckpoints: %v", err) - } - if len(sessions) != 2 { - t.Fatalf("sessions len = %d, want 2", len(sessions)) - } - if filepath.Base(sessions[0].Path) != "session-new.json" { - t.Fatalf("first session = %s, want session-new.json", sessions[0].Path) - } - if sessions[0].Messages != 1 || sessions[0].Model != "new" { - t.Fatalf("session metadata = %+v", sessions[0]) - } -} - -func writeSessionFile(t *testing.T, path string, data CheckpointData) { - t.Helper() - raw, err := json.Marshal(data) - if err != nil { - t.Fatalf("marshal session: %v", err) - } - if err := os.WriteFile(path, raw, 0o644); err != nil { - t.Fatalf("write session: %v", err) - } -} - -func TestSanitizeMessagesForSave(t *testing.T) { - msgs := []*aop.Message{ - { - Role: "assistant", - Content: []*aop.Content{ - aop.Reasoning("thinking..."), - aop.Text("part1"), - aop.Image("image/png", []byte("binary-image-data")), - aop.Text("part2"), - }, - }, - } - out := sanitizeMessagesForSave(msgs) - if len(out) != 1 { - t.Fatalf("len = %d", len(out)) - } - if got := provider.MessageText(out[0]); got != "part1part2" { - t.Errorf("content = %q, want %q", got, "part1part2") - } - for _, part := range out[0].Content { - if part.GetMedia() != nil { - t.Error("media parts should be stripped after sanitize") - } - } - if got := provider.MessageReasoning(out[0]); got != "thinking..." { - t.Errorf("reasoning = %q, want preserved", got) - } -} - -func TestLoadCheckpointNotFound(t *testing.T) { - _, err := LoadCheckpoint("/nonexistent/path.json") - if err == nil { - t.Error("expected error for nonexistent file") - } -} diff --git a/agent/loop.go b/agent/loop.go index 0381c6c0..e01b5385 100644 --- a/agent/loop.go +++ b/agent/loop.go @@ -12,7 +12,6 @@ import ( "github.com/chainreactors/aiscan/agent/inbox" "github.com/chainreactors/aiscan/agent/provider" aop "github.com/chainreactors/aiscan/aop" - "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/core/tool" "github.com/chainreactors/aiscan/core/truncate" @@ -49,9 +48,6 @@ func runLoop(ctx context.Context, cfg Config) (*Result, error) { em.errorEvt(result.Err, isRetryableError(result.Err)) } emitRunEnd(ctx, cfg, result) - if cfg.OnRunEnd != nil { - cfg.OnRunEnd(result) - } } return result, err } @@ -521,7 +517,9 @@ type toolExecution struct { func runToolCall(ctx context.Context, cfg Config, assistantMsg *aop.Message, tc *aop.ToolCall, turn int) toolExecution { startedAt := time.Now() - toolCtx := output.ContextWithCallID(ctx, tc.Id) + toolCtx := tool.ContextWithInvocation(ctx, tool.Invocation{ + CallID: tc.Id, SessionID: cfg.SessionID, TurnID: cfg.TurnID, Emitter: cfg.AgentName, + }) toolCtx = withToolAgentConfig(toolCtx, cfg) toolCtx = inbox.ContextWithInbox(toolCtx, cfg.Inbox) execution := beforeToolCall(toolCtx, cfg, assistantMsg, tc) diff --git a/agent/subagent.go b/agent/subagent.go index 85b029cf..844d6f05 100644 --- a/agent/subagent.go +++ b/agent/subagent.go @@ -13,7 +13,6 @@ import ( "github.com/chainreactors/aiscan/agent/inbox" "github.com/chainreactors/aiscan/agent/provider" aop "github.com/chainreactors/aiscan/aop" - "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/core/tool" types "github.com/chainreactors/aiscan/pkg/types" @@ -136,7 +135,7 @@ func (t *SubAgentTool) create(ctx context.Context, prompt, typeName, name, mode, if err != nil { return "", err } - parentToolCallID := output.CallIDFromContext(ctx) + parentToolCallID := tool.InvocationFromContext(ctx).CallID if parentToolCallID == "" { return "", fmt.Errorf("subagent create requires the spawning tool call id") } diff --git a/agent/subagent_test.go b/agent/subagent_test.go index 0dedcbcc..04e04957 100644 --- a/agent/subagent_test.go +++ b/agent/subagent_test.go @@ -9,7 +9,6 @@ import ( "github.com/chainreactors/aiscan/agent/inbox" aop "github.com/chainreactors/aiscan/aop" "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/core/output" coretool "github.com/chainreactors/aiscan/core/tool" "github.com/chainreactors/aiscan/pkg/commands" types "github.com/chainreactors/aiscan/pkg/types" @@ -24,7 +23,7 @@ func TestSubAgentSyncReturnsResult(t *testing.T) { }) tool := NewSubAgentTool(nil) - ctx := output.ContextWithCallID(withToolAgentConfig(context.Background(), parent.Cfg), "spawn-sync") + ctx := coretool.ContextWithInvocation(withToolAgentConfig(context.Background(), parent.Cfg), coretool.Invocation{CallID: "spawn-sync"}) result, err := tool.Execute(ctx, `{"action":"create","mode":"sync","name":"worker","prompt":"do the work"}`) if err != nil { t.Fatalf("Execute() error = %v", err) @@ -83,7 +82,7 @@ func TestSubAgentUsesExecutingAgentContext(t *testing.T) { Bus: bus, }) - ctx := output.ContextWithCallID(withToolAgentConfig(context.Background(), active.Cfg), "spawn-context") + ctx := coretool.ContextWithInvocation(withToolAgentConfig(context.Background(), active.Cfg), coretool.Invocation{CallID: "spawn-context"}) if _, err := tool.Execute(ctx, `{"action":"create","mode":"async","name":"context-worker","prompt":"work"}`); err != nil { t.Fatalf("Execute() error = %v", err) } diff --git a/agent/types.go b/agent/types.go index ddfa2181..aa3172f3 100644 --- a/agent/types.go +++ b/agent/types.go @@ -132,13 +132,10 @@ type Config struct { TokenBudget int Logger telemetry.Logger TransformContext TransformContextFunc - Bus EventEmitter + Bus aop.EventEmitter // Hooks is the typed extension registry shared by a runtime and its derived // agents. Nil means no handlers and keeps the dispatch fast path allocation-free. - Hooks *hooks.Registry - // OnRunEnd fires once per run with the final result — replaces the old - // EventAgentEnd Messages subscription for session persistence. - OnRunEnd func(*Result) + Hooks *hooks.Registry BeforeToolCall func(context.Context, BeforeToolCallContext) (*BeforeToolCallResult, error) AfterToolCall func(context.Context, AfterToolCallContext) (*AfterToolCallResult, error) MaxTurns int @@ -167,21 +164,21 @@ type Config struct { // Builder methods — each returns a modified copy (Config is a value type). -func (c Config) WithProvider(p Provider) Config { c.Provider = p; return c } -func (c Config) WithTools(t tool.Executor) Config { c.Tools = t; return c } -func (c Config) WithModel(m string) Config { c.Model = m; return c } -func (c Config) WithSystemPrompt(s string) Config { c.SystemPrompt = s; return c } -func (c Config) WithMessages(msgs []*aop.Message) Config { c.Messages = msgs; return c } -func (c Config) WithStream(s bool) Config { c.Stream = s; return c } -func (c Config) WithInbox(ib inbox.Inbox) Config { c.Inbox = ib; return c } -func (c Config) WithLogger(l telemetry.Logger) Config { c.Logger = l; return c } -func (c Config) WithBus(b EventEmitter) Config { c.Bus = b; return c } -func (c Config) WithMaxTokens(n int) Config { c.MaxTokens = n; return c } -func (c Config) WithContextWindow(n int) Config { c.ContextWindow = n; return c } -func (c Config) WithTemperature(t float64) Config { c.Temperature = &t; return c } -func (c Config) WithMaxRetries(n int) Config { c.MaxRetries = n; return c } -func (c Config) WithTokenBudget(n int) Config { c.TokenBudget = n; return c } -func (c Config) WithExpander(e *inbox.Expander) Config { c.Expander = e; return c } +func (c Config) WithProvider(p Provider) Config { c.Provider = p; return c } +func (c Config) WithTools(t tool.Executor) Config { c.Tools = t; return c } +func (c Config) WithModel(m string) Config { c.Model = m; return c } +func (c Config) WithSystemPrompt(s string) Config { c.SystemPrompt = s; return c } +func (c Config) WithMessages(msgs []*aop.Message) Config { c.Messages = msgs; return c } +func (c Config) WithStream(s bool) Config { c.Stream = s; return c } +func (c Config) WithInbox(ib inbox.Inbox) Config { c.Inbox = ib; return c } +func (c Config) WithLogger(l telemetry.Logger) Config { c.Logger = l; return c } +func (c Config) WithBus(b aop.EventEmitter) Config { c.Bus = b; return c } +func (c Config) WithMaxTokens(n int) Config { c.MaxTokens = n; return c } +func (c Config) WithContextWindow(n int) Config { c.ContextWindow = n; return c } +func (c Config) WithTemperature(t float64) Config { c.Temperature = &t; return c } +func (c Config) WithMaxRetries(n int) Config { c.MaxRetries = n; return c } +func (c Config) WithTokenBudget(n int) Config { c.TokenBudget = n; return c } +func (c Config) WithExpander(e *inbox.Expander) Config { c.Expander = e; return c } func (c Config) WithTransformContext(fn TransformContextFunc) Config { c.TransformContext = fn return c @@ -191,7 +188,6 @@ func (c Config) WithSessionID(id string) Config { c.SessionID = id; func (c Config) WithTurnID(id string) Config { c.TurnID = id; return c } func (c Config) WithAgentName(name string) Config { c.AgentName = name; return c } func (c Config) WithHooks(r *hooks.Registry) Config { c.Hooks = r; return c } -func (c Config) WithOnRunEnd(fn func(*Result)) Config { c.OnRunEnd = fn; return c } func (c Config) WithLoopScheduler(s *LoopScheduler) Config { c.LoopScheduler = s return c diff --git a/aop/helpers.go b/aop/helpers.go index 72620cfe..fd1c39c0 100644 --- a/aop/helpers.go +++ b/aop/helpers.go @@ -8,6 +8,10 @@ import ( "google.golang.org/protobuf/types/known/anypb" ) +type EventEmitter interface { + Emit(*Event) +} + const JSONMediaType = "application/json" func JSONValue(value any) (*EncodedValue, error) { diff --git a/aop/tool/artifact.go b/aop/tool/artifact.go new file mode 100644 index 00000000..dd11aebb --- /dev/null +++ b/aop/tool/artifact.go @@ -0,0 +1,10 @@ +package tool + +const ( + ArtifactKindService = "service" + ArtifactKindWeb = "web" + ArtifactKindWeakpass = "weakpass" + ArtifactKindVuln = "vuln" + ArtifactKindSummary = "summary" + ArtifactKindError = "error" +) diff --git a/aop/tool/protocol.pb.go b/aop/tool/protocol.pb.go index 6a613764..cc730fa0 100644 --- a/aop/tool/protocol.pb.go +++ b/aop/tool/protocol.pb.go @@ -89,6 +89,7 @@ type Progress struct { Target string `protobuf:"bytes,3,opt,name=target,proto3" json:"target,omitempty"` Timestamp *timestamppb.Timestamp `protobuf:"bytes,5,opt,name=timestamp,proto3" json:"timestamp,omitempty"` Text string `protobuf:"bytes,6,opt,name=text,proto3" json:"text,omitempty"` + CallId string `protobuf:"bytes,7,opt,name=call_id,json=callId,proto3" json:"call_id,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -151,6 +152,13 @@ func (x *Progress) GetText() string { return "" } +func (x *Progress) GetCallId() string { + if x != nil { + return x.CallId + } + return "" +} + // Artifact carries one scanner-native structured record. Nodes remain thin: // only the server normalizes these records into canonical SCO documents. type Artifact struct { @@ -161,6 +169,7 @@ type Artifact struct { Data []byte `protobuf:"bytes,4,opt,name=data,proto3" json:"data,omitempty"` MediaType string `protobuf:"bytes,5,opt,name=media_type,json=mediaType,proto3" json:"media_type,omitempty"` Timestamp *timestamppb.Timestamp `protobuf:"bytes,6,opt,name=timestamp,proto3" json:"timestamp,omitempty"` + CallId string `protobuf:"bytes,7,opt,name=call_id,json=callId,proto3" json:"call_id,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -237,6 +246,13 @@ func (x *Artifact) GetTimestamp() *timestamppb.Timestamp { return nil } +func (x *Artifact) GetCallId() string { + if x != nil { + return x.CallId + } + return "" +} + type ProtocolMessage struct { state protoimpl.MessageState `protogen:"open.v1"` // Types that are valid to be assigned to Message: @@ -344,12 +360,13 @@ const file_aop_tool_protocol_proto_rawDesc = "" + "\n" + "session_id\x18\x01 \x01(\tR\tsessionId\x12\x17\n" + "\aturn_id\x18\x02 \x01(\tR\x06turnId\x12!\n" + - "\x04call\x18\x03 \x01(\v2\r.aop.ToolCallR\x04call\"\x90\x01\n" + + "\x04call\x18\x03 \x01(\v2\r.aop.ToolCallR\x04call\"\xa9\x01\n" + "\bProgress\x12\x12\n" + "\x04tool\x18\x01 \x01(\tR\x04tool\x12\x16\n" + "\x06target\x18\x03 \x01(\tR\x06target\x128\n" + "\ttimestamp\x18\x05 \x01(\v2\x1a.google.protobuf.TimestampR\ttimestamp\x12\x12\n" + - "\x04text\x18\x06 \x01(\tR\x04textJ\x04\b\x02\x10\x03J\x04\b\x04\x10\x05\"\xb7\x01\n" + + "\x04text\x18\x06 \x01(\tR\x04text\x12\x17\n" + + "\acall_id\x18\a \x01(\tR\x06callIdJ\x04\b\x02\x10\x03J\x04\b\x04\x10\x05\"\xd0\x01\n" + "\bArtifact\x12\x12\n" + "\x04tool\x18\x01 \x01(\tR\x04tool\x12\x12\n" + "\x04kind\x18\x02 \x01(\tR\x04kind\x12\x16\n" + @@ -357,7 +374,8 @@ const file_aop_tool_protocol_proto_rawDesc = "" + "\x04data\x18\x04 \x01(\fR\x04data\x12\x1d\n" + "\n" + "media_type\x18\x05 \x01(\tR\tmediaType\x128\n" + - "\ttimestamp\x18\x06 \x01(\v2\x1a.google.protobuf.TimestampR\ttimestamp\"\xa6\x01\n" + + "\ttimestamp\x18\x06 \x01(\v2\x1a.google.protobuf.TimestampR\ttimestamp\x12\x17\n" + + "\acall_id\x18\a \x01(\tR\x06callId\"\xa6\x01\n" + "\x0fProtocolMessage\x120\n" + "\bprogress\x18\n" + " \x01(\v2\x12.aop.tool.ProgressH\x00R\bprogress\x12$\n" + diff --git a/cmd/aiscan/cli.go b/cmd/aiscan/cli.go index 7f22d057..e1985559 100644 --- a/cmd/aiscan/cli.go +++ b/cmd/aiscan/cli.go @@ -125,7 +125,7 @@ func aiscan() { return } if option.ViewFile != "" { - if err := output.RenderFile(option.ViewFile, option.ViewFormat, option.ViewOutput); err != nil { + if err := output.RenderEventFile(option.ViewFile, option.ViewFormat, option.OutputFile); err != nil { fmt.Fprintf(os.Stderr, "error: %s\n", err) os.Exit(1) } @@ -231,6 +231,9 @@ func parseCLI(args []string) (parsedCLI, error) { if cli.Timeout > 0 { option.Timeout = cli.Timeout } + if err := validateSessionFileFlags(option); err != nil { + return parsedCLI{}, err + } if mode == cfg.RunModeNoCommand { return parsedCLI{Option: option, Mode: cfg.RunModeNoCommand}, nil @@ -301,10 +304,18 @@ func parseScannerCLI(scannerName string, rootArgs, scannerRest []string) (parsed if err != nil { return parsedCLI{}, err } + } else { + scannerArgs, err = applyScannerPersistenceArgs(scannerRest, &option) + if err != nil { + return parsedCLI{}, err + } } if boolFlagEnabled(scannerArgs, "--debug") { option.Debug = true } + if err := validateSessionFileFlags(option); err != nil { + return parsedCLI{}, err + } return parsedCLI{ Option: option, Mode: cfg.RunModeScanner, @@ -312,7 +323,42 @@ func parseScannerCLI(scannerName string, rootArgs, scannerRest []string) (parsed }, nil } +func validateSessionFileFlags(option cfg.Option) error { + if strings.TrimSpace(option.Resume) != "" && strings.TrimSpace(option.OutputFile) != "" { + return fmt.Errorf("--resume/-r and --file/-f are mutually exclusive") + } + return nil +} + +func applyScannerPersistenceArgs(args []string, option *cfg.Option) ([]string, error) { + out := make([]string, 0, len(args)) + for i := 0; i < len(args); i++ { + arg := args[i] + key, value, hasValue := strings.Cut(arg, "=") + switch key { + case "--file", "-f": + resolved, err := flagValue(arg, hasValue, value, args, &i) + if err != nil { + return nil, err + } + option.OutputFile = resolved + case "--resume", "-r": + resolved, err := flagValue(arg, hasValue, value, args, &i) + if err != nil { + return nil, err + } + option.Resume = resolved + case "--save-session": + option.SaveSession = !hasValue || truthyFlagValue(value) + default: + out = append(out, arg) + } + } + return out, nil +} + func mergeManualScannerOptions(option *cfg.Option, manual cfg.Option) { + option.OutputFile = cfg.ResolveString(manual.OutputFile, option.OutputFile) option.Provider = cfg.ResolveString(manual.Provider, option.Provider) option.BaseURL = cfg.ResolveString(manual.BaseURL, option.BaseURL) option.APIKey = cfg.ResolveString(manual.APIKey, option.APIKey) @@ -345,6 +391,10 @@ func mergeManualScannerOptions(option *cfg.Option, manual cfg.Option) { option.Prompt = cfg.ResolveString(manual.Prompt, option.Prompt) option.TaskFile = cfg.ResolveString(manual.TaskFile, option.TaskFile) option.WebURL = cfg.ResolveString(manual.WebURL, option.WebURL) + option.Resume = cfg.ResolveString(manual.Resume, option.Resume) + if manual.SaveSession { + option.SaveSession = true + } if len(manual.Skills) > 0 { option.Skills = append(option.Skills, manual.Skills...) } @@ -519,7 +569,9 @@ var scannerKnownFlags = []knownFlag{ } }}, {names: []string{"--resume"}, arity: 1, apply: func(o *cfg.Option, v string) { o.Resume = v }}, + {names: []string{"-r"}, arity: 1, apply: func(o *cfg.Option, v string) { o.Resume = v }}, {names: []string{"--save-session"}, arity: 0, apply: func(o *cfg.Option, _ string) { o.SaveSession = true }}, + {names: []string{"--file", "-f"}, arity: 1, apply: func(o *cfg.Option, v string) { o.OutputFile = v }}, } var rootOnlyFlagValueArity = map[string]int{ diff --git a/cmd/aiscan/cli_test.go b/cmd/aiscan/cli_test.go index 3c1ee2fa..37d0fe97 100644 --- a/cmd/aiscan/cli_test.go +++ b/cmd/aiscan/cli_test.go @@ -127,6 +127,60 @@ func TestParseCLIScannerKeepsToolTimeoutAfterCommand(t *testing.T) { } } +func TestParseCLIExtractsUnifiedFileForAgentAndScanners(t *testing.T) { + tests := []struct { + args []string + wantArgs []string + }{ + {args: []string{"agent", "-p", "hello", "-f", "agent.jsonl"}}, + {args: []string{"scan", "-i", "127.0.0.1", "-f", "scan.jsonl"}, wantArgs: []string{"scan", "-i", "127.0.0.1"}}, + {args: []string{"gogo", "-i", "127.0.0.1", "-p", "80", "-f", "gogo.jsonl"}, wantArgs: []string{"gogo", "-i", "127.0.0.1", "-p", "80"}}, + } + for _, test := range tests { + t.Run(test.args[0], func(t *testing.T) { + parsed, err := parseCLI(test.args) + if err != nil { + t.Fatalf("parseCLI: %v", err) + } + wantFile := test.args[len(test.args)-1] + if parsed.Option.OutputFile != wantFile { + t.Fatalf("output file = %q, want %q", parsed.Option.OutputFile, wantFile) + } + if test.wantArgs != nil && !reflect.DeepEqual(parsed.ScannerArgs, test.wantArgs) { + t.Fatalf("scanner args = %#v, want %#v", parsed.ScannerArgs, test.wantArgs) + } + }) + } +} + +func TestParseCLIViewUsesUnifiedInputAndFileFlags(t *testing.T) { + parsed, err := parseCLI([]string{"-F", "session.jsonl", "-o", "markdown", "-f", "session.md"}) + if err != nil { + t.Fatalf("parseCLI: %v", err) + } + if parsed.Option.ViewFile != "session.jsonl" || parsed.Option.ViewFormat != "markdown" || parsed.Option.OutputFile != "session.md" { + t.Fatalf("view options = %#v", parsed.Option.MiscOptions) + } +} + +func TestParseCLIRejectsResumeWithExplicitFile(t *testing.T) { + for _, args := range [][]string{ + {"agent", "-r", "session.jsonl", "-f", "other.jsonl"}, + {"scan", "-i", "127.0.0.1", "-r", "session.jsonl", "-f", "other.jsonl"}, + } { + if _, err := parseCLI(args); err == nil || !strings.Contains(err.Error(), "mutually exclusive") { + t.Fatalf("parseCLI(%v) error = %v", args, err) + } + } + parsed, err := parseCLI([]string{"agent", "-r", "session.jsonl", "--save-session"}) + if err != nil { + t.Fatalf("resume + save-session: %v", err) + } + if parsed.Option.Resume != "session.jsonl" || !parsed.Option.SaveSession { + t.Fatalf("resume options = %#v", parsed.Option) + } +} + func TestParseCLIRootTimeoutAppliesToAgent(t *testing.T) { parsed, err := parseCLI([]string{"--timeout", "45", "agent", "-p", "test"}) if err != nil { diff --git a/cmd/aiscan/setup.go b/cmd/aiscan/setup.go index 0990d6cc..c62adae6 100644 --- a/cmd/aiscan/setup.go +++ b/cmd/aiscan/setup.go @@ -11,8 +11,6 @@ import ( aop "github.com/chainreactors/aiscan/aop" "github.com/chainreactors/aiscan/core/capability" cfg "github.com/chainreactors/aiscan/core/config" - "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/pidlock" "github.com/chainreactors/aiscan/core/resources" "github.com/chainreactors/aiscan/core/telemetry" @@ -42,7 +40,7 @@ func scannerInit(ctx context.Context, a *runner.App, rc runner.ApplicationConfig es := initEngines(ctx, rc.Scanner, logger) a.Engines = es registerScannerCommands(a.Commands, es, rc.Scanner, rc.Tools, - a.Provider, a.ProviderConfig, a.Skills, a.DataBus, logger) + a.Provider, a.ProviderConfig, a.Skills, a.Events, logger) } func initEngines(ctx context.Context, sc runner.ScannerConfig, logger telemetry.Logger) *engine.Set { @@ -69,7 +67,7 @@ func initEngines(ctx context.Context, sc runner.ScannerConfig, logger telemetry. return engineSet } -func registerScannerCommands(cmdReg *commands.CommandRegistry, engineSet *engine.Set, scanCfg runner.ScannerConfig, toolCfg runner.ToolConfig, llmProvider agent.Provider, providerConfig agent.ProviderConfig, skillStore *skills.Store, dataBus *eventbus.Bus[output.ToolDataEvent], logger telemetry.Logger) { +func registerScannerCommands(cmdReg *commands.CommandRegistry, engineSet *engine.Set, scanCfg runner.ScannerConfig, toolCfg runner.ToolConfig, llmProvider agent.Provider, providerConfig agent.ProviderConfig, skillStore *skills.Store, agentEvents aop.EventEmitter, logger telemetry.Logger) { var scanOpts []scan.Option if scanCfg.AIEnabled && llmProvider != nil { scannerParent := agent.NewAgent(agent.Config{ @@ -79,6 +77,7 @@ func registerScannerCommands(cmdReg *commands.CommandRegistry, engineSet *engine MaxTokens: providerConfig.MaxTokens, ContextWindow: providerConfig.ContextWindow, Logger: logger, + Bus: agentEvents, }) scanOpts = append(scanOpts, scan.WithParent(scannerParent)) scanOpts = append(scanOpts, scan.WithDeepBrowserFunc(func(ctx context.Context, targetURL string) (string, error) { @@ -105,7 +104,7 @@ func registerScannerCommands(cmdReg *commands.CommandRegistry, engineSet *engine Logger: logger, TavilyKeys: toolCfg.TavilyKeys, PlaywrightSession: toolCfg.PlaywrightSession, - DataBus: dataBus, + Events: agentEvents, } commands.Provide(deps, scan.OptsKey, scanOpts) if engineSet != nil { diff --git a/cmd/aiscan/web_full.go b/cmd/aiscan/web_full.go index 5b6312af..c37906ea 100644 --- a/cmd/aiscan/web_full.go +++ b/cmd/aiscan/web_full.go @@ -15,8 +15,9 @@ import ( "sync" "time" + aop "github.com/chainreactors/aiscan/aop" + toolpb "github.com/chainreactors/aiscan/aop/tool" cfg "github.com/chainreactors/aiscan/core/config" - "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/telemetry" node "github.com/chainreactors/aiscan/pkg/node" "github.com/chainreactors/aiscan/pkg/runner" @@ -174,11 +175,17 @@ func runWeb(ctx context.Context, option, explicitOption *cfg.Option, opts webCom } func wireWebApp(application *runner.App, ingestor webservice.ArtifactIngestor) { - if application == nil || ingestor == nil || application.Artifacts == nil { + if application == nil || ingestor == nil || application.EventBus == nil { return } - application.Artifacts.SetHandler(func(artifact output.ToolArtifact) { - _ = ingestor.IngestArtifact(context.Background(), artifact.CallID, artifact) + application.EventBus.Subscribe(func(event *aop.Event) { + if event == nil || event.GetExtension() == nil { + return + } + artifact := new(toolpb.Artifact) + if event.GetExtension().MessageIs(artifact) && event.GetExtension().UnmarshalTo(artifact) == nil { + _ = ingestor.IngestArtifact(context.Background(), artifact) + } }) } diff --git a/cmd/aiscan/web_full_test.go b/cmd/aiscan/web_full_test.go index 8f007181..9bc28896 100644 --- a/cmd/aiscan/web_full_test.go +++ b/cmd/aiscan/web_full_test.go @@ -9,11 +9,14 @@ import ( "runtime" "testing" + aop "github.com/chainreactors/aiscan/aop" + 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/output" "github.com/chainreactors/aiscan/pkg/runner" types "github.com/chainreactors/aiscan/pkg/types" + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/types/known/anypb" ) func TestWebConfigStoreStagesBeforeAtomicCommit(t *testing.T) { @@ -70,28 +73,29 @@ func TestWebConfigStoreStagesBeforeAtomicCommit(t *testing.T) { } func TestWireWebAppBindsRawArtifactsForReloadedApp(t *testing.T) { - bus := eventbus.New[output.ToolDataEvent]() - application := &runner.App{Artifacts: output.NewArtifactStream(bus)} - defer application.Artifacts.Close() + bus := eventbus.New[*aop.Event]() + application := &runner.App{EventBus: bus} ingestor := &recordingArtifactIngestor{} wireWebApp(application, ingestor) - bus.Emit(output.ToolDataEvent{ - Tool: "gogo", Kind: output.ToolDataService, CallID: "scan-1", - Data: map[string]string{"ip": "127.0.0.1"}, + extension, err := anypb.New(&toolpb.Artifact{ + Tool: "gogo", Kind: toolpb.ArtifactKindService, CallId: "scan-1", Data: []byte(`{"ip":"127.0.0.1"}`), }) - if ingestor.operationID != "scan-1" || ingestor.artifact.Tool != "gogo" { + if err != nil { + t.Fatal(err) + } + bus.Emit(&aop.Event{SessionId: "session-1", Payload: &aop.Event_Extension{Extension: extension}}) + if ingestor.artifact == nil || ingestor.artifact.CallId != "scan-1" || ingestor.artifact.Tool != "gogo" { t.Fatalf("artifact was not forwarded: %+v", ingestor.artifact) } } type recordingArtifactIngestor struct { - operationID string - artifact output.ToolArtifact + artifact *toolpb.Artifact } -func (i *recordingArtifactIngestor) IngestArtifact(_ context.Context, operationID string, artifact output.ToolArtifact) error { - i.operationID, i.artifact = operationID, artifact +func (i *recordingArtifactIngestor) IngestArtifact(_ context.Context, artifact *toolpb.Artifact) error { + i.artifact = proto.CloneOf(artifact) return nil } diff --git a/core/config/options.go b/core/config/options.go index ab93e21c..fcf88abc 100644 --- a/core/config/options.go +++ b/core/config/options.go @@ -17,7 +17,7 @@ type Option struct { AgentOptions `group:"Agent Options" config:"agent"` IOAOptions `group:"Server Options" config:"ioa"` ReconOptions `group:"Recon Options" config:"recon"` - OutputOptions `group:"Agent Output Options" config:"output"` + OutputOptions `group:"Output Options" config:"output"` MiscOptions `group:"Miscellaneous Options" config:"misc"` ScanConfig ScanConfigOptions `no-flag:"true" config:"scan"` SearchConfig SearchConfigOptions `no-flag:"true" config:"search"` @@ -86,8 +86,8 @@ type AgentOptions struct { ServerURL string `long:"server-url" config:"server_url" description:"AIScan Web server URL for AOP, remote REPL and PTY access"` WebURL string `long:"web-url" config:"web_url" description:"Deprecated alias for --server-url" hidden:"true"` Transport string `long:"transport" config:"transport" description:"Agent transport: auto, local, web, or stdio" default:"auto"` - Resume string `long:"resume" description:"Resume session from a saved session file path"` - SaveSession bool `long:"save-session" config:"save_session" description:"Auto-save conversation to .aiscan/sessions/ after each agent run (default: off)"` + Resume string `short:"r" long:"resume" description:"Resume agent context from an AOP JSONL session file"` + SaveSession bool `long:"save-session" config:"save_session" description:"Auto-select a .aiscan/sessions/*.jsonl recording path"` CaptureProviderFrames bool `long:"capture-provider-frames" config:"capture_provider_frames" description:"Emit exact provider request/response frames as sensitive AOP events"` } @@ -139,9 +139,9 @@ type MiscOptions struct { ConfigFile string `short:"c" long:"config" description:"Path to config file (default: ./aiscan.yaml, /aiscan.yaml)"` DataDir string `long:"data-dir" config:"data_dir" description:"Data directory for cache, arsenal, history (default: /.aiscan)"` InitConfig bool `long:"init" description:"Generate default aiscan.yaml and exit"` - ViewFile string `short:"F" long:"view" description:"View a scan record JSONL file"` + ViewFile string `short:"F" long:"view" description:"View an AOP event JSONL file"` ViewFormat string `short:"o" long:"output" description:"Output format for -F: terminal (default), markdown" default:"terminal"` - ViewOutput string `short:"f" long:"file" description:"Write -F output to file instead of stdout"` + OutputFile string `short:"f" long:"file" description:"Write scan and agent events to a streaming JSONL file (with -F: write rendered output)"` Debug bool `long:"debug" config:"debug" description:"Enable debug logging"` Verbose []bool `short:"v" long:"verbose" description:"Increase verbosity (-v thinking and tool previews, -vv full tool results)"` Quiet bool `short:"q" long:"quiet" config:"quiet" description:"Quiet mode — only show final result"` diff --git a/core/output/artifact_stream.go b/core/output/artifact_stream.go deleted file mode 100644 index 83de3d67..00000000 --- a/core/output/artifact_stream.go +++ /dev/null @@ -1,83 +0,0 @@ -package output - -import ( - "encoding/json" - "sync" - "time" - - "github.com/chainreactors/aiscan/core/eventbus" -) - -// ToolArtifact is one scanner-native structured record. Nodes forward it -// unchanged; only the server normalizes it into canonical SCO documents. -type ToolArtifact struct { - Tool string - Kind string - Target string - Data json.RawMessage - CallID string - Timestamp time.Time -} - -// ArtifactStream serializes structured ToolData events and exposes one -// lifecycle-bound handler for either a node transport or the local server. -type ArtifactStream struct { - mu sync.RWMutex - unsub func() - handler func(ToolArtifact) -} - -func NewArtifactStream(bus *eventbus.Bus[ToolDataEvent]) *ArtifactStream { - stream := &ArtifactStream{} - if bus != nil { - stream.unsub = bus.Subscribe(stream.handle) - } - return stream -} - -func (s *ArtifactStream) handle(event ToolDataEvent) { - if event.Kind == ToolDataProgress || event.Data == nil { - return - } - data, err := json.Marshal(event.Data) - if err != nil { - return - } - s.mu.RLock() - handler := s.handler - s.mu.RUnlock() - if handler == nil { - return - } - handler(ToolArtifact{ - Tool: event.Tool, - Kind: event.Kind, - Target: event.Target, - Data: data, - CallID: event.CallID, - Timestamp: event.Timestamp, - }) -} - -func (s *ArtifactStream) SetHandler(handler func(ToolArtifact)) { - if s == nil { - return - } - s.mu.Lock() - s.handler = handler - s.mu.Unlock() -} - -func (s *ArtifactStream) Close() { - if s == nil { - return - } - s.mu.Lock() - unsub := s.unsub - s.unsub = nil - s.handler = nil - s.mu.Unlock() - if unsub != nil { - unsub() - } -} diff --git a/core/output/artifact_stream_test.go b/core/output/artifact_stream_test.go deleted file mode 100644 index cc222ac3..00000000 --- a/core/output/artifact_stream_test.go +++ /dev/null @@ -1,49 +0,0 @@ -package output - -import ( - "encoding/json" - "testing" - "time" - - "github.com/chainreactors/aiscan/core/eventbus" -) - -func TestArtifactStreamForwardsStructuredRecords(t *testing.T) { - bus := eventbus.New[ToolDataEvent]() - stream := NewArtifactStream(bus) - defer stream.Close() - - var got ToolArtifact - stream.SetHandler(func(artifact ToolArtifact) { got = artifact }) - timestamp := time.Now() - bus.Emit(ToolDataEvent{ - Tool: "gogo", Kind: ToolDataService, Target: "192.0.2.1:80", - Data: map[string]any{"ip": "192.0.2.1", "port": "80"}, - CallID: "scan-1", Timestamp: timestamp, - }) - - if got.Tool != "gogo" || got.Kind != ToolDataService || got.CallID != "scan-1" { - t.Fatalf("unexpected artifact metadata: %+v", got) - } - var data map[string]string - if err := json.Unmarshal(got.Data, &data); err != nil { - t.Fatal(err) - } - if data["ip"] != "192.0.2.1" || data["port"] != "80" { - t.Fatalf("unexpected artifact data: %v", data) - } -} - -func TestArtifactStreamIgnoresProgressAndStopsOnClose(t *testing.T) { - bus := eventbus.New[ToolDataEvent]() - stream := NewArtifactStream(bus) - count := 0 - stream.SetHandler(func(ToolArtifact) { count++ }) - - bus.Emit(ToolDataEvent{Kind: ToolDataProgress, Data: "running"}) - stream.Close() - bus.Emit(ToolDataEvent{Tool: "gogo", Kind: ToolDataService, Data: map[string]string{"ip": "192.0.2.1"}}) - if count != 0 { - t.Fatalf("handler called %d times", count) - } -} diff --git a/core/output/jsonl.go b/core/output/jsonl.go new file mode 100644 index 00000000..b00bcf2f --- /dev/null +++ b/core/output/jsonl.go @@ -0,0 +1,188 @@ +package output + +import ( + "bufio" + "fmt" + "io" + "os" + "path/filepath" + "strings" + "sync" + + aop "github.com/chainreactors/aiscan/aop" + "github.com/chainreactors/aiscan/core/eventbus" + "google.golang.org/protobuf/encoding/protojson" +) + +// 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. +func ScanJSONL(path string, visit func(*aop.Event) error) error { + file, err := os.Open(path) + if err != nil { + return fmt.Errorf("open AOP JSONL: %w", err) + } + defer file.Close() + 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] != '{' { + continue + } + 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 visit != nil { + if err := visit(event); err != nil { + return err + } + } + } + if err := scanner.Err(); err != nil { + return fmt.Errorf("read AOP JSONL: %w", err) + } + return nil +} + +func ReadJSONL(path string) ([]*aop.Event, error) { + var events []*aop.Event + err := ScanJSONL(path, func(event *aop.Event) error { + events = append(events, event) + return nil + }) + return events, err +} + +// 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 +} + +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{} + if err := recorder.Switch(path); err != nil { + return nil, err + } + recorder.unsub = bus.Subscribe(func(event *aop.Event) { + if err := recorder.Write(event); err != nil { + recorder.mu.Lock() + if recorder.err == nil { + recorder.err = err + } + recorder.mu.Unlock() + } + }) + return recorder, nil +} + +func openJSONL(path string) (*os.File, string, error) { + path = filepath.Clean(strings.TrimSpace(path)) + if path == "." || path == "" { + return nil, "", fmt.Errorf("AOP JSONL path is required") + } + if dir := filepath.Dir(path); dir != "." && dir != "" { + if err := os.MkdirAll(dir, 0o755); err != nil { + return nil, "", fmt.Errorf("create AOP JSONL directory: %w", err) + } + } + file, err := os.OpenFile(path, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o644) + if err != nil { + return nil, "", fmt.Errorf("open AOP JSONL %s: %w", path, err) + } + return file, path, nil +} + +func ValidateJSONLTarget(path string) error { + file, _, err := openJSONL(path) + if err != nil { + return err + } + return file.Close() +} + +func (r *JSONLRecorder) Switch(path string) error { + file, clean, err := openJSONL(path) + if err != nil { + return err + } + r.mu.Lock() + old := r.file + r.file = file + r.path = clean + r.mu.Unlock() + 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 (r *JSONLRecorder) Path() string { + if r == nil { + return "" + } + r.mu.Lock() + defer r.mu.Unlock() + return r.path +} + +func (r *JSONLRecorder) Write(event *aop.Event) error { + if r == nil || event == nil { + return nil + } + line, err := (protojson.MarshalOptions{UseProtoNames: true}).Marshal(event) + if err != nil { + return fmt.Errorf("marshal AOP JSONL event: %w", err) + } + line = append(line, '\n') + r.mu.Lock() + defer r.mu.Unlock() + if r.file == nil { + return io.ErrClosedPipe + } + n, err := r.file.Write(line) + if err == nil && n != len(line) { + err = io.ErrShortWrite + } + return err +} + +func (r *JSONLRecorder) Close() error { + if r == nil { + return nil + } + r.mu.Lock() + unsub := r.unsub + r.unsub = nil + file := r.file + recordErr := r.err + r.file = nil + r.path = "" + r.err = nil + r.mu.Unlock() + if unsub != nil { + unsub() + } + if file != nil { + if err := file.Close(); recordErr == nil { + recordErr = err + } + } + return recordErr +} diff --git a/core/output/jsonl_test.go b/core/output/jsonl_test.go new file mode 100644 index 00000000..fcfc0fe6 --- /dev/null +++ b/core/output/jsonl_test.go @@ -0,0 +1,83 @@ +package output + +import ( + "fmt" + "path/filepath" + "sync" + "testing" + + aop "github.com/chainreactors/aiscan/aop" + "github.com/chainreactors/aiscan/core/eventbus" +) + +func TestJSONLRecorderWritesConcurrentEventsAsCompleteLines(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) + } + const count = 64 + var wg sync.WaitGroup + for i := 0; i < count; i++ { + wg.Add(1) + go func(index int) { + defer wg.Done() + bus.Emit(&aop.Event{ + Id: fmt.Sprintf("event-%d", index), SessionId: "session", Emitter: "test", + Payload: &aop.Event_Message{Message: &aop.Message{Role: "assistant", Content: []*aop.Content{aop.Text(fmt.Sprint(index))}}}, + }) + }(i) + } + wg.Wait() + if err := recorder.Close(); err != nil { + t.Fatal(err) + } + events, err := ReadJSONL(path) + if err != nil { + t.Fatal(err) + } + if len(events) != count { + t.Fatalf("events = %d, want %d", len(events), count) + } +} + +func TestJSONLRecorderSwitchesFilesWithoutReplayingHistory(t *testing.T) { + dir := t.TempDir() + first := filepath.Join(dir, "first.jsonl") + second := filepath.Join(dir, "second.jsonl") + bus := eventbus.New[*aop.Event]() + recorder, err := NewJSONLRecorder(bus, first) + if err != nil { + t.Fatal(err) + } + bus.Emit(jsonlTestMessage("first")) + if err := recorder.Switch(second); err != nil { + t.Fatal(err) + } + bus.Emit(jsonlTestMessage("second")) + if err := recorder.Close(); err != nil { + t.Fatal(err) + } + firstEvents, err := ReadJSONL(first) + if err != nil { + t.Fatal(err) + } + secondEvents, err := ReadJSONL(second) + if err != nil { + t.Fatal(err) + } + if len(firstEvents) != 1 || firstEvents[0].Id != "first" { + t.Fatalf("first events = %#v", firstEvents) + } + if len(secondEvents) != 1 || secondEvents[0].Id != "second" { + t.Fatalf("second events = %#v", secondEvents) + } +} + +func jsonlTestMessage(id string) *aop.Event { + return &aop.Event{ + Id: id, SessionId: "session", Emitter: "test", + Payload: &aop.Event_Message{Message: &aop.Message{Role: "assistant", Content: []*aop.Content{aop.Text(id)}}}, + } +} diff --git a/core/output/timeline.go b/core/output/render.go similarity index 68% rename from core/output/timeline.go rename to core/output/render.go index 04f17bbf..154fb183 100644 --- a/core/output/timeline.go +++ b/core/output/render.go @@ -1,7 +1,6 @@ package output import ( - "bufio" "encoding/json" "fmt" "io" @@ -14,74 +13,24 @@ import ( types "github.com/chainreactors/aiscan/pkg/types" "github.com/charmbracelet/glamour" "github.com/muesli/termenv" - "google.golang.org/protobuf/encoding/protojson" ) -// --------------------------------------------------------------------------- -// Core types -// --------------------------------------------------------------------------- - -type TimelineEntry struct { - Timestamp time.Time - Type string - Data any -} - -// --------------------------------------------------------------------------- -// Parse -// --------------------------------------------------------------------------- - -func ParseTimelineFile(path string) ([]TimelineEntry, error) { - f, err := os.Open(path) - if err != nil { - return nil, err - } - defer f.Close() - - var entries []TimelineEntry - scanner := bufio.NewScanner(f) - scanner.Buffer(make([]byte, 0, 256*1024), 10*1024*1024) - for scanner.Scan() { - line := scanner.Bytes() - if len(line) == 0 || line[0] != '{' { - continue - } - if e, ok := parseLine(line); ok { - entries = append(entries, e) - } - } - return entries, scanner.Err() -} - -func parseLine(line []byte) (TimelineEntry, bool) { - event := new(aop.Event) - if protojson.Unmarshal(line, event) == nil && event.SessionId != "" && event.Payload != nil { - timestamp := time.Time{} - if event.EmittedAt != nil { - timestamp = event.EmittedAt.AsTime() - } - return TimelineEntry{Timestamp: timestamp, Type: aop.Kind(event), Data: event}, true - } - return TimelineEntry{}, false -} - // --------------------------------------------------------------------------- // Render entry points // --------------------------------------------------------------------------- -func RenderTimeline(w io.Writer, entries []TimelineEntry) error { - _, err := io.WriteString(w, renderMD(BuildTimelineMarkdown(entries))) +func RenderEvents(w io.Writer, events []*aop.Event) error { + _, err := io.WriteString(w, renderMD(BuildEventMarkdown(events))) return err } -func RenderTimelineMarkdown(w io.Writer, entries []TimelineEntry) error { - _, err := io.WriteString(w, BuildTimelineMarkdown(entries)) +func RenderEventsMarkdown(w io.Writer, events []*aop.Event) error { + _, err := io.WriteString(w, BuildEventMarkdown(events)) return err } -// RenderFile renders an AOP Event ProtoJSONL file. Raw scanner JSONL files are -// intentionally not accepted as agent timelines. -func RenderFile(path, format, outputPath string) error { +// RenderEventFile renders an AOP Event ProtoJSONL file. +func RenderEventFile(path, format, outputPath string) error { var writer io.Writer = os.Stdout if outputPath != "" { file, err := os.Create(outputPath) @@ -91,26 +40,30 @@ func RenderFile(path, format, outputPath string) error { defer file.Close() writer = file } - entries, err := ParseTimelineFile(path) + events, err := ReadJSONL(path) if err != nil { return err } if strings.EqualFold(format, "markdown") || strings.EqualFold(format, "md") { - return RenderTimelineMarkdown(writer, entries) + return RenderEventsMarkdown(writer, events) } - return RenderTimeline(writer, entries) + return RenderEvents(writer, events) } -func BuildTimelineMarkdown(entries []TimelineEntry) string { +func BuildEventMarkdown(events []*aop.Event) string { var sb strings.Builder - sess := collectSessionMeta(entries) - writeHeader(&sb, &sess) + sessions := collectSessionMeta(events) + writtenHeaders := make(map[string]bool) - for _, e := range entries { - switch d := e.Data.(type) { - case *aop.Event: - writeAOPMarkdown(&sb, d) + for _, event := range events { + if event != nil && event.SessionId != "" && !writtenHeaders[event.SessionId] { + if sb.Len() > 0 { + sb.WriteString("\n") + } + writeHeader(&sb, sessions[event.SessionId]) + writtenHeaders[event.SessionId] = true } + writeAOPMarkdown(&sb, event) } return sb.String() } @@ -161,39 +114,45 @@ func (s *sessionMeta) duration() time.Duration { return s.endTS.Sub(s.startTS) } -func collectSessionMeta(entries []TimelineEntry) sessionMeta { - var m sessionMeta - for _, e := range entries { - switch d := e.Data.(type) { - case *aop.Event: - if m.id == "" { - m.id = d.SessionId +func collectSessionMeta(events []*aop.Event) map[string]*sessionMeta { + sessions := make(map[string]*sessionMeta) + for _, event := range events { + if event == nil || event.SessionId == "" { + continue + } + m := sessions[event.SessionId] + if m == nil { + m = &sessionMeta{id: event.SessionId} + sessions[event.SessionId] = m + } + timestamp := time.Time{} + if event.EmittedAt != nil { + timestamp = event.EmittedAt.AsTime() + } + switch payload := event.Payload.(type) { + case *aop.Event_SessionStarted: + m.startTS = timestamp + m.parentID = payload.SessionStarted.ParentSessionId + if payload.SessionStarted.Model != "" && m.model == "" { + m.model = payload.SessionStarted.Model + } + case *aop.Event_SessionEnded: + m.endTS = timestamp + case *aop.Event_TurnStarted: + m.turns++ + case *aop.Event_TurnEnded: + m.endTS = timestamp + m.stop = payload.TurnEnded.StopReason + if payload.TurnEnded.Usage != nil && payload.TurnEnded.Usage.TotalTokens > 0 { + m.totalTokens = int(payload.TurnEnded.Usage.TotalTokens) } - switch payload := d.Payload.(type) { - case *aop.Event_SessionStarted: - m.startTS = e.Timestamp - m.parentID = payload.SessionStarted.ParentSessionId - if payload.SessionStarted.Model != "" && m.model == "" { - m.model = payload.SessionStarted.Model - } - case *aop.Event_SessionEnded: - m.endTS = e.Timestamp - case *aop.Event_TurnStarted: - m.turns++ - case *aop.Event_TurnEnded: - m.endTS = e.Timestamp - m.stop = payload.TurnEnded.StopReason - if payload.TurnEnded.Usage != nil && payload.TurnEnded.Usage.TotalTokens > 0 { - m.totalTokens = int(payload.TurnEnded.Usage.TotalTokens) - } - case *aop.Event_Usage: - if payload.Usage.TotalTokens > 0 { - m.totalTokens = int(payload.Usage.TotalTokens) - } + case *aop.Event_Usage: + if payload.Usage.TotalTokens > 0 { + m.totalTokens = int(payload.Usage.TotalTokens) } } } - return m + return sessions } // --------------------------------------------------------------------------- @@ -201,25 +160,25 @@ func collectSessionMeta(entries []TimelineEntry) sessionMeta { // --------------------------------------------------------------------------- var ( - timelineRenderer *glamour.TermRenderer - timelineRendererErr error - timelineRendererOnce sync.Once + eventRenderer *glamour.TermRenderer + eventRendererErr error + eventRendererOnce sync.Once ) -func getTimelineRenderer() (*glamour.TermRenderer, error) { - timelineRendererOnce.Do(func() { - timelineRenderer, timelineRendererErr = glamour.NewTermRenderer( +func getEventRenderer() (*glamour.TermRenderer, error) { + eventRendererOnce.Do(func() { + eventRenderer, eventRendererErr = glamour.NewTermRenderer( glamour.WithAutoStyle(), glamour.WithColorProfile(termenv.ANSI), glamour.WithEmoji(), glamour.WithWordWrap(120), ) }) - return timelineRenderer, timelineRendererErr + return eventRenderer, eventRendererErr } func renderMD(md string) string { - r, err := getTimelineRenderer() + r, err := getEventRenderer() if err != nil { return md } diff --git a/core/output/render_test.go b/core/output/render_test.go new file mode 100644 index 00000000..889f57f5 --- /dev/null +++ b/core/output/render_test.go @@ -0,0 +1,136 @@ +package output + +import ( + "os" + "path/filepath" + "strings" + "testing" + "time" + + aop "github.com/chainreactors/aiscan/aop" + toolpb "github.com/chainreactors/aiscan/aop/tool" + "github.com/chainreactors/aiscan/core/eventbus" + 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" +) + +func TestParseLineReadsNativeAOPEnvelope(t *testing.T) { + event := renderEvent(&aop.Event{Payload: &aop.Event_Message{Message: &aop.Message{ + Id: "m-1", Role: "assistant", Content: []*aop.Content{aop.Text("hello")}, + }}}) + raw, err := protojson.Marshal(event) + if err != nil { + t.Fatal(err) + } + + parsed := new(aop.Event) + if err := protojson.Unmarshal(raw, parsed); err != nil { + t.Fatal(err) + } + if markdown := BuildEventMarkdown([]*aop.Event{parsed}); !strings.Contains(markdown, "hello") { + t.Fatalf("event markdown = %q", markdown) + } +} + +func TestEventRendererRendersStructuredToolResult(t *testing.T) { + event := renderEvent(&aop.Event{Payload: &aop.Event_ToolResult{ToolResult: &aop.ToolResult{ + CallId: "call-1", Name: "scan", Output: []*aop.Content{ + aop.Text("three ports"), aop.Image("image/png", []byte("x")), + }, + }}}) + markdown := BuildEventMarkdown([]*aop.Event{event}) + if !strings.Contains(markdown, "three ports") { + t.Fatalf("event markdown = %q", markdown) + } +} + +func TestEventRendererFormatsPreformattedCommandAtPresentationBoundary(t *testing.T) { + event := renderEvent(&aop.Event{Payload: &aop.Event_Message{Message: &aop.Message{ + Id: "command-1", Role: "assistant", Content: []*aop.Content{aop.Text("one\ntwo")}, + }}}) + _ = types.SetCommandDetail(event, &types.CommandDetail{Line: "/status", Presentation: "preformatted"}) + markdown := BuildEventMarkdown([]*aop.Event{event}) + if !strings.Contains(markdown, "```\none\ntwo\n```") { + t.Fatalf("event markdown = %q", markdown) + } +} + +func TestEventRendererDoesNotRenderStructuredArtifactPayloads(t *testing.T) { + extension, err := anypb.New(&toolpb.Artifact{ + Tool: "gogo", Kind: toolpb.ArtifactKindService, Target: "127.0.0.1:443", Data: []byte(`{"secret":"structured-only"}`), + }) + if err != nil { + t.Fatal(err) + } + event := renderEvent(&aop.Event{Payload: &aop.Event_Extension{Extension: extension}}) + markdown := BuildEventMarkdown([]*aop.Event{event}) + if strings.Contains(markdown, "structured-only") || strings.Contains(markdown, "127.0.0.1:443") { + t.Fatalf("artifact leaked into generic markdown: %q", markdown) + } +} + +func TestRenderEventFileFormatsTheSameAOPJSONLStream(t *testing.T) { + dir := t.TempDir() + inputPath := filepath.Join(dir, "session.jsonl") + outputPath := filepath.Join(dir, "session.md") + bus := eventbus.New[*aop.Event]() + writer, err := NewJSONLRecorder(bus, inputPath) + if err != nil { + t.Fatal(err) + } + events := []*aop.Event{ + renderEvent(&aop.Event{Payload: &aop.Event_SessionStarted{SessionStarted: &aop.SessionStarted{Model: "test-model"}}}), + renderEvent(&aop.Event{Payload: &aop.Event_Message{Message: &aop.Message{Id: "m-1", Role: "user", Content: []*aop.Content{aop.Text("rendered prompt")}}}}), + renderEvent(&aop.Event{Payload: &aop.Event_Message{Message: &aop.Message{Id: "m-2", Role: "assistant", Content: []*aop.Content{aop.Text("rendered answer")}}}}), + } + for _, event := range events { + bus.Emit(event) + } + if err := writer.Close(); err != nil { + t.Fatal(err) + } + if err := RenderEventFile(inputPath, "markdown", outputPath); err != nil { + t.Fatalf("RenderEventFile: %v", err) + } + rendered, err := os.ReadFile(outputPath) + if err != nil { + t.Fatal(err) + } + text := string(rendered) + for _, expected := range []string{"test-model", "rendered prompt", "rendered answer"} { + if !strings.Contains(text, expected) { + t.Fatalf("formatted output missing %q:\n%s", expected, text) + } + } +} + +func TestEventRendererSeparatesContinuationSessionHeaders(t *testing.T) { + root := renderEvent(&aop.Event{Payload: &aop.Event_SessionStarted{SessionStarted: &aop.SessionStarted{Model: "test-model"}}}) + root.SessionId = "root-1" + continuation := renderEvent(&aop.Event{Payload: &aop.Event_SessionStarted{SessionStarted: &aop.SessionStarted{ + Model: "test-model", ParentSessionId: "root-1", + }}}) + continuation.SessionId = "cont-2" + + markdown := BuildEventMarkdown([]*aop.Event{root, continuation}) + if strings.Count(markdown, "# Agent ") != 2 { + t.Fatalf("session headers = %q", markdown) + } + if !strings.Contains(markdown, "# Agent `root-1`") || !strings.Contains(markdown, "# Agent `cont-2 ← root-1`") { + t.Fatalf("continuation headers = %q", markdown) + } + if strings.Contains(markdown, "# Agent `root-1 ← root-1`") { + t.Fatalf("root header inherited continuation parent: %q", markdown) + } +} + +func renderEvent(event *aop.Event) *aop.Event { + event.Id = "e-1" + event.SessionId = "session-1" + event.TurnId = "turn-1" + event.Emitter = "aiscan" + event.EmittedAt = timestamppb.New(time.Date(2026, 7, 20, 0, 0, 0, 0, time.UTC)) + return event +} diff --git a/core/output/timeline_test.go b/core/output/timeline_test.go deleted file mode 100644 index c3b1c088..00000000 --- a/core/output/timeline_test.go +++ /dev/null @@ -1,65 +0,0 @@ -package output - -import ( - "strings" - "testing" - "time" - - aop "github.com/chainreactors/aiscan/aop" - types "github.com/chainreactors/aiscan/pkg/types" - "google.golang.org/protobuf/encoding/protojson" - "google.golang.org/protobuf/types/known/timestamppb" -) - -func TestParseLineReadsNativeAOPEnvelope(t *testing.T) { - event := timelineEvent(&aop.Event{Payload: &aop.Event_Message{Message: &aop.Message{ - Id: "m-1", Role: "assistant", Content: []*aop.Content{aop.Text("hello")}, - }}}) - raw, err := protojson.Marshal(event) - if err != nil { - t.Fatal(err) - } - - entry, ok := parseLine(raw) - if !ok { - t.Fatal("native AOP event was not parsed") - } - if _, ok := entry.Data.(*aop.Event); !ok { - t.Fatalf("entry data type = %T", entry.Data) - } - if markdown := BuildTimelineMarkdown([]TimelineEntry{entry}); !strings.Contains(markdown, "hello") { - t.Fatalf("timeline markdown = %q", markdown) - } -} - -func TestTimelineRendersStructuredToolResult(t *testing.T) { - event := timelineEvent(&aop.Event{Payload: &aop.Event_ToolResult{ToolResult: &aop.ToolResult{ - CallId: "call-1", Name: "scan", Output: []*aop.Content{ - aop.Text("three ports"), aop.Image("image/png", []byte("x")), - }, - }}}) - markdown := BuildTimelineMarkdown([]TimelineEntry{{Timestamp: event.EmittedAt.AsTime(), Type: aop.Kind(event), Data: event}}) - if !strings.Contains(markdown, "three ports") { - t.Fatalf("timeline markdown = %q", markdown) - } -} - -func TestTimelineFormatsPreformattedCommandAtPresentationBoundary(t *testing.T) { - event := timelineEvent(&aop.Event{Payload: &aop.Event_Message{Message: &aop.Message{ - Id: "command-1", Role: "assistant", Content: []*aop.Content{aop.Text("one\ntwo")}, - }}}) - _ = types.SetCommandDetail(event, &types.CommandDetail{Line: "/status", Presentation: "preformatted"}) - markdown := BuildTimelineMarkdown([]TimelineEntry{{Timestamp: event.EmittedAt.AsTime(), Type: aop.Kind(event), Data: event}}) - if !strings.Contains(markdown, "```\none\ntwo\n```") { - t.Fatalf("timeline markdown = %q", markdown) - } -} - -func timelineEvent(event *aop.Event) *aop.Event { - event.Id = "e-1" - event.SessionId = "session-1" - event.TurnId = "turn-1" - event.Emitter = "aiscan" - event.EmittedAt = timestamppb.New(time.Date(2026, 7, 20, 0, 0, 0, 0, time.UTC)) - return event -} diff --git a/core/output/tool_data.go b/core/output/tool_data.go deleted file mode 100644 index f42355bd..00000000 --- a/core/output/tool_data.go +++ /dev/null @@ -1,38 +0,0 @@ -package output - -import ( - "context" - "time" -) - -type toolCallIDKey struct{} - -func ContextWithCallID(ctx context.Context, callID string) context.Context { - return context.WithValue(ctx, toolCallIDKey{}, callID) -} - -func CallIDFromContext(ctx context.Context) string { - if v, ok := ctx.Value(toolCallIDKey{}).(string); ok { - return v - } - return "" -} - -type ToolDataEvent struct { - Tool string `json:"tool"` - Kind string `json:"kind"` - Target string `json:"target,omitempty"` - Data any `json:"data"` - CallID string `json:"call_id,omitempty"` - Timestamp time.Time `json:"timestamp"` -} - -const ( - ToolDataService = "service" - ToolDataWeb = "web" - ToolDataWeakpass = "weakpass" - ToolDataVuln = "vuln" - // ToolDataProgress carries one raw stdout/stderr line of a foreground - // command, streamed while it runs. - ToolDataProgress = "progress" -) diff --git a/core/output/types.go b/core/output/types.go index 241b1668..e9495808 100644 --- a/core/output/types.go +++ b/core/output/types.go @@ -7,7 +7,7 @@ import ( ) // ScanResult is private collector state. Scanner-native records leave a node -// only as ToolArtifact messages; the server owns normalization and persistence. +// only as canonical aop.tool.Artifact messages. type ScanResult struct { Summary Summary `json:"summary"` GOGO []*parsers.GOGOResult `json:"gogo,omitempty"` diff --git a/core/tool/context.go b/core/tool/context.go index 6f3eca8b..5092a73e 100644 --- a/core/tool/context.go +++ b/core/tool/context.go @@ -7,7 +7,11 @@ type invocationContextKey struct{} // Invocation carries executor-owned context that must not become model-facing // tool arguments. type Invocation struct { - WorkDir string + WorkDir string + CallID string + SessionID string + TurnID string + Emitter string } func ContextWithInvocation(ctx context.Context, invocation Invocation) context.Context { diff --git a/docs/mechanisms.md b/docs/mechanisms.md index f3a8dac9..2ccd4e84 100644 --- a/docs/mechanisms.md +++ b/docs/mechanisms.md @@ -243,9 +243,11 @@ AIScan 产品事件使用 AOP core 的 typed Any 插槽;例如 scan 完成通 ### 命令展示边界 -跨界面 Runtime 命令通过 typed AOP command detail 标记 `presentation: preformatted`。Web timeline 在最终展示边界生成自适应 Markdown code fence;Runtime、Session 和 transport 不再处理 Markdown 或终端格式。 +跨界面 Runtime 命令通过 typed AOP command detail 标记 `presentation: preformatted`。Web 展示层和 `-F` 格式化入口只在最终展示边界生成自适应 Markdown code fence;Runtime、Session 和 transport 不处理 Markdown 或终端格式。 -**文件**: `pkg/tui/banner.go`, `pkg/tui/commands.go`, `pkg/types/extensions.go`, `core/output/timeline.go` +Session 持久化只有一条路径:所有需要持久化的 agent、scan 和 tool artifact 都先成为 `aop.Event`,经同一个 EventBus 流式追加到 ProtoJSONL。`-r` 从该文件恢复上下文并继续追加;`/resume` 关闭旧 session 后切换到目标文件;`/clear` 和 `/compact` 仅在当前文件内创建 continuation session。Progress 只用于实时传输,不持久化,也不存在 checkpoint、snapshot 或 timeline replay 文件机制。 + +**文件**: `pkg/tui/banner.go`, `pkg/tui/commands.go`, `pkg/types/extensions.go`, `core/output/jsonl.go`, `core/output/render.go`, `pkg/runner/session_jsonl.go` --- diff --git a/docs/protocol-architecture.md b/docs/protocol-architecture.md index 7ee5d4d2..c99eec92 100644 --- a/docs/protocol-architecture.md +++ b/docs/protocol-architecture.md @@ -135,8 +135,8 @@ AOP 应用面只额外暴露一个双向流服务: - Session 和 Scan 以 protobuf 为存储真相; - AOP 历史只存 `aop.Event` ProtoJSON; -- Scanner 文件输出只写 libcstx SCO JSONL; -- 不保留旧扁平 DTO 列、Record/Timeline 双写或 fallback read。 +- CLI `-f` 将 agent、scan 和 scanner-native artifact 统一写入同一个 append-only `aop.Event` ProtoJSONL; +- `-r`、`/resume` 和 `-F` 直接读取该事件流,不保留 checkpoint/snapshot 文件、Record/Timeline 双写或 replay/fallback 管线。 历史读取是纯查询,不派发 Agent frame、不收敛 operation,也不复制 terminal event。 diff --git a/docs/scan.md b/docs/scan.md index 9d188bb9..d4f9b4c3 100644 --- a/docs/scan.md +++ b/docs/scan.md @@ -121,8 +121,8 @@ scan 提供 `quick` 和 `full` 两种预设模式,通过 `--mode` 参数选择 | `--deep` | 对发现的 Web 资产进行 AI 动态测试 | | | `-j, --json` | JSON Lines 输出 | | | `--report` | Markdown 报告输出 | | -| `-f, --file` | 输出写入文件(自动去除 ANSI 颜色) | | -| `-F, --view` | 回放之前保存的 JSONL 扫描记录 | | +| `-f, --file` | 将 scan、agent 与结构化 tool artifact 事件流式追加为 AOP ProtoJSONL | | +| `-F, --view` | 按需格式化同一份 AOP ProtoJSONL(`-o terminal\|markdown`,可用 `-f` 指定展示输出文件) | | | `--trace` | 显示内部 pipeline 事件流(调试用) | | | `--no-color` | 禁用终端颜色 | | | `--debug` | 启用 trace + 底层扫描器 debug 日志 | | diff --git a/examples/acp/server/main.go b/examples/acp/server/main.go index 30f1d5fb..68832414 100644 --- a/examples/acp/server/main.go +++ b/examples/acp/server/main.go @@ -13,8 +13,9 @@ import ( "syscall" "time" + aop "github.com/chainreactors/aiscan/aop" + toolpb "github.com/chainreactors/aiscan/aop/tool" cfg "github.com/chainreactors/aiscan/core/config" - "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/telemetry" node "github.com/chainreactors/aiscan/pkg/node" "github.com/chainreactors/aiscan/pkg/runner" @@ -30,9 +31,15 @@ func newHeadlessHandler(store *webservice.SQLiteStore, app *runner.App, ingestor pool := webservice.NewAgentPool(service.Hub()) pool.SetArtifactIngestor(ingestor) service.SetAgentPool(pool) - if app != nil && app.Artifacts != nil && ingestor != nil { - app.Artifacts.SetHandler(func(artifact output.ToolArtifact) { - _ = ingestor.IngestArtifact(context.Background(), artifact.CallID, artifact) + if app != nil && app.EventBus != nil && ingestor != nil { + app.EventBus.Subscribe(func(event *aop.Event) { + if event == nil || event.GetExtension() == nil { + return + } + artifact := new(toolpb.Artifact) + if event.GetExtension().MessageIs(artifact) && event.GetExtension().UnmarshalTo(artifact) == nil { + _ = ingestor.IngestArtifact(context.Background(), artifact) + } }) } return service, pool, web.NewHandler(service, nil, nil) diff --git a/pkg/commands/factory.go b/pkg/commands/factory.go index 6606ff74..ef1df105 100644 --- a/pkg/commands/factory.go +++ b/pkg/commands/factory.go @@ -5,10 +5,9 @@ import ( "github.com/chainreactors/aiscan/agent/hooks" "github.com/chainreactors/aiscan/agent/provider" + aop "github.com/chainreactors/aiscan/aop" "github.com/chainreactors/aiscan/core/capability" "github.com/chainreactors/aiscan/core/deps" - "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/telemetry" ) @@ -42,7 +41,7 @@ type Deps struct { NodeMeta map[string]any TavilyKeys string // comma-separated Tavily API keys PlaywrightSession string - DataBus *eventbus.Bus[output.ToolDataEvent] + Events aop.EventEmitter Hooks *hooks.Registry } diff --git a/pkg/node/agent.go b/pkg/node/agent.go index 08b02296..d12f564b 100644 --- a/pkg/node/agent.go +++ b/pkg/node/agent.go @@ -65,8 +65,7 @@ func runRemoteAgent(ctx context.Context, option *cfg.Option, logger telemetry.Lo Name: runner.ResolveIOANodeName(option), Registry: application.Commands, AgentSubscribe: rt.Subscribe, - DataBus: application.DataBus, - Artifacts: application.Artifacts, + Progress: application.Progress, Logger: logger, Chat: chatHandler, AgentRuntime: rt, diff --git a/pkg/node/connection.go b/pkg/node/connection.go index 5e0bb951..a8a472f6 100644 --- a/pkg/node/connection.go +++ b/pkg/node/connection.go @@ -4,8 +4,8 @@ import ( "context" aop "github.com/chainreactors/aiscan/aop" + toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/telemetry" coreterminal "github.com/chainreactors/aiscan/core/terminal" "github.com/chainreactors/aiscan/pkg/commands" @@ -27,8 +27,7 @@ type connectionConfig struct { JSONFrames bool Registry *commands.CommandRegistry AgentSubscribe func(func(*aop.Event)) func() - DataBus *eventbus.Bus[output.ToolDataEvent] - Artifacts *output.ArtifactStream + Progress *eventbus.Bus[*toolpb.Progress] Logger telemetry.Logger Chat *chatAgentHandler // AgentRuntime handles the AOP core/command namespaces directly via diff --git a/pkg/node/proto_connection.go b/pkg/node/proto_connection.go index fd18310a..eb5293d4 100644 --- a/pkg/node/proto_connection.go +++ b/pkg/node/proto_connection.go @@ -222,12 +222,17 @@ func serveAgentConnection(ctx context.Context, cc connectionConfig, logger telem replyTo := "" if event.GetToolResult() != nil { replyTo = event.GetToolResult().GetCallId() + } else if extension := event.GetExtension(); extension != nil { + artifact := new(toolpb.Artifact) + if extension.MessageIs(artifact) && extension.UnmarshalTo(artifact) == nil { + replyTo = artifact.CallId + } } send(replyTo, &aop.ProtocolMessage{Message: &aop.ProtocolMessage_Event{Event: event}}) }) defer unsubscribe() } - if detach := attachToolEvents(cc.DataBus, cc.Artifacts, send); detach != nil { + if detach := attachToolProgress(cc.Progress, send); detach != nil { defer detach() } if cc.Status != nil { @@ -435,15 +440,27 @@ func handleAgentToolMessage(ctx context.Context, cc connectionConfig, envelope * return } operationID := envelope.GetId() + 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)}, + }) + } taskCtx, taskCancel := context.WithCancel(ctx) trackOperation(operationsMu, operations, operationID, taskCancel) go func() { defer finishOperation(operationsMu, operations, operationID, taskCancel) - event, err := runner.ExecuteToolRequest(taskCtx, operationID, request, cc.Registry, cc.DataBus) + event, err := runner.ExecuteToolRequest(taskCtx, operationID, request, cc.Registry, cc.Progress) if err != nil { fail(err.Error()) return } + if cc.AgentRuntime != nil { + cc.AgentRuntime.EmitEvent(event) + } send(replyTo, &aop.ProtocolMessage{Message: &aop.ProtocolMessage_Event{Event: event}}) }() } diff --git a/pkg/node/toolnode.go b/pkg/node/toolnode.go index 4d5c142d..3b4e0af4 100644 --- a/pkg/node/toolnode.go +++ b/pkg/node/toolnode.go @@ -6,12 +6,10 @@ import ( "os" "runtime" "strings" - "time" aop "github.com/chainreactors/aiscan/aop" toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/pkg/runner" @@ -19,7 +17,6 @@ import ( "github.com/chainreactors/aiscan/skills" protobuf "google.golang.org/protobuf/proto" "google.golang.org/protobuf/types/known/structpb" - "google.golang.org/protobuf/types/known/timestamppb" ) // ToolNodeConfig configures a tool-only runner: an outbound WebSocket @@ -30,13 +27,13 @@ type ToolNodeConfig struct { ServerURL string WSPath string // ID is the stable node identity used by Cairn as the runner primary key. - ID string - Token string - Registry *commands.CommandRegistry - DataBus *eventbus.Bus[output.ToolDataEvent] - Artifacts *output.ArtifactStream - Logger telemetry.Logger - Version string + ID string + Token string + Registry *commands.CommandRegistry + Events *eventbus.Bus[*aop.Event] + Progress *eventbus.Bus[*toolpb.Progress] + Logger telemetry.Logger + Version string // JSONFrames switches the hub wire to standard ProtoJSON text frames; // hubs expecting binary protobuf (AIScan) leave it false. JSONFrames bool @@ -77,67 +74,40 @@ func RunToolNode(ctx context.Context, cfg ToolNodeConfig) error { skillStore, _ := skills.LoadEmbeddedStore() menu = func() []*types.CommandSpec { return runner.RegistryCommandCatalog(cfg.Registry, skillStore) } } - artifacts := cfg.Artifacts - if artifacts == nil && cfg.DataBus != nil { - artifacts = output.NewArtifactStream(cfg.DataBus) - defer artifacts.Close() + 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, - DataBus: cfg.DataBus, - Artifacts: artifacts, - Logger: logger, - NodeID: runnerID, - Runtime: runnerRuntime, - Capabilities: []string{"pty", "file", "exec", "tool", "artifact"}, - Menu: menu, - RunnerFileRPC: true, - JSONFrames: cfg.JSONFrames, + ServerURL: cfg.ServerURL, + WSPath: cfg.WSPath, + Name: runnerID, + Token: cfg.Token, + Registry: cfg.Registry, + AgentSubscribe: subscribe, + Progress: cfg.Progress, + Logger: logger, + NodeID: runnerID, + Runtime: runnerRuntime, + Capabilities: []string{"pty", "file", "exec", "tool", "artifact"}, + Menu: menu, + RunnerFileRPC: true, + JSONFrames: cfg.JSONFrames, }) } -// attachToolEvents forwards progress and scanner-native artifacts onto the hub -// connection. Nodes never normalize artifacts into SCO. -func attachToolEvents(dataBus *eventbus.Bus[output.ToolDataEvent], artifacts *output.ArtifactStream, send func(string, protobuf.Message)) func() { - if dataBus == nil && artifacts == nil { +// attachToolProgress forwards ephemeral progress onto the tool protocol. Raw +// artifacts travel as canonical AOP extension events through AgentSubscribe. +func attachToolProgress(progressBus *eventbus.Bus[*toolpb.Progress], send func(string, protobuf.Message)) func() { + if progressBus == nil { return nil } - var unsub func() - if dataBus != nil { - unsub = dataBus.Subscribe(func(event output.ToolDataEvent) { - if event.Kind != output.ToolDataProgress { - return - } - text, ok := event.Data.(string) - if !ok || text == "" { - return - } - timestamp := event.Timestamp - if timestamp.IsZero() { - timestamp = time.Now() - } - send(event.CallID, &toolpb.ProtocolMessage{Message: &toolpb.ProtocolMessage_Progress{Progress: &toolpb.Progress{ - Tool: event.Tool, Target: event.Target, Text: text, Timestamp: timestamppb.New(timestamp), - }}}) - }) - } - if artifacts != nil { - artifacts.SetHandler(func(artifact output.ToolArtifact) { - timestamp := artifact.Timestamp - if timestamp.IsZero() { - timestamp = time.Now() - } - send(artifact.CallID, &toolpb.ProtocolMessage{Message: &toolpb.ProtocolMessage_Artifact{Artifact: &toolpb.Artifact{ - Tool: artifact.Tool, Kind: artifact.Kind, Target: artifact.Target, - Data: append([]byte(nil), artifact.Data...), MediaType: aop.JSONMediaType, - Timestamp: timestamppb.New(timestamp), - }}}) - }) - } + unsub := progressBus.Subscribe(func(progress *toolpb.Progress) { + if progress == nil || progress.Text == "" { + return + } + send(progress.CallId, &toolpb.ProtocolMessage{Message: &toolpb.ProtocolMessage_Progress{Progress: protobuf.CloneOf(progress)}}) + }) var once bool return func() { if once { @@ -147,8 +117,5 @@ func attachToolEvents(dataBus *eventbus.Bus[output.ToolDataEvent], artifacts *ou if unsub != nil { unsub() } - if artifacts != nil { - artifacts.SetHandler(nil) - } } } diff --git a/pkg/node/toolnode_test.go b/pkg/node/toolnode_test.go index 08ba99c0..64b191e6 100644 --- a/pkg/node/toolnode_test.go +++ b/pkg/node/toolnode_test.go @@ -13,13 +13,13 @@ import ( filepb "github.com/chainreactors/aiscan/aop/file" toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/tool" "github.com/chainreactors/aiscan/pkg/commands" types "github.com/chainreactors/aiscan/pkg/types" "github.com/gorilla/websocket" "google.golang.org/protobuf/encoding/protojson" protobuf "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/types/known/anypb" ) var testUpgrader = websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} @@ -127,9 +127,16 @@ func (h *hubScript) serveHTTP(w http.ResponseWriter, r *http.Request) { } switch value := message.(type) { case *aop.ProtocolMessage: - if result := value.GetEvent().GetToolResult(); result != nil { + event := value.GetEvent() + if result := event.GetToolResult(); result != nil { h.toolResult <- result } + if extension := event.GetExtension(); extension != nil { + artifact := new(toolpb.Artifact) + if extension.MessageIs(artifact) && extension.UnmarshalTo(artifact) == nil { + h.artifact <- artifact + } + } case *toolpb.ProtocolMessage: if progress := value.GetProgress(); progress != nil { h.progress <- progress.Text @@ -179,7 +186,8 @@ func TestRunToolNodeWireInterop(t *testing.T) { Name: "gogo", Usage: "gogo [OPTIONS]", DescriptionPath: "aiscan://skills/aiscan/okf/easm/gogo.md", }, "scanner") - dataBus := eventbus.New[output.ToolDataEvent]() + events := eventbus.New[*aop.Event]() + progress := eventbus.New[*toolpb.Progress]() hub := newHubScript(t) server := httptest.NewServer(http.HandlerFunc(hub.serveHTTP)) defer server.Close() @@ -187,7 +195,7 @@ func TestRunToolNodeWireInterop(t *testing.T) { defer cancel() errCh := make(chan error, 1) go func() { - errCh <- RunToolNode(ctx, ToolNodeConfig{ServerURL: server.URL, WSPath: "/ws/runner", ID: "runner-1", Token: "test-token", Registry: registry, DataBus: dataBus, Version: "test"}) + errCh <- RunToolNode(ctx, ToolNodeConfig{ServerURL: server.URL, WSPath: "/ws/runner", ID: "runner-1", Token: "test-token", Registry: registry, Events: events, Progress: progress, Version: "test"}) }() hello := wait(t, hub.registered, "hello") if hello.Name != "runner-1" || hello.NodeId != "runner-1" { @@ -217,12 +225,16 @@ func TestRunToolNodeWireInterop(t *testing.T) { if got := catalog.Commands[0].GetDescription(); got != "Use this playbook when working with gogo for host, port, service, banner, fingerprint, or vulnerability-hint discovery." { t.Fatalf("gogo description = %q", got) } - dataBus.Emit(output.ToolDataEvent{ - Tool: "gogo", Kind: output.ToolDataService, Target: "192.0.2.1:80", - Data: map[string]string{"ip": "192.0.2.1", "port": "80"}, CallID: "exec-1", + extension, err := anypb.New(&toolpb.Artifact{ + Tool: "gogo", Kind: toolpb.ArtifactKindService, Target: "192.0.2.1:80", + Data: []byte(`{"ip":"192.0.2.1","port":"80"}`), CallId: "exec-1", MediaType: aop.JSONMediaType, }) + if err != nil { + t.Fatal(err) + } + events.Emit(&aop.Event{SessionId: "session-1", TurnId: "turn-1", Emitter: "gogo", Payload: &aop.Event_Extension{Extension: extension}}) artifact := wait(t, hub.artifact, "tool artifact") - if artifact.Tool != "gogo" || artifact.Kind != output.ToolDataService || string(artifact.Data) != `{"ip":"192.0.2.1","port":"80"}` { + if artifact.Tool != "gogo" || artifact.Kind != toolpb.ArtifactKindService || artifact.CallId != "exec-1" || string(artifact.Data) != `{"ip":"192.0.2.1","port":"80"}` { t.Fatalf("artifact = %+v data=%s", artifact, artifact.Data) } if line := wait(t, hub.progress, "tool progress"); line != "streamed" { diff --git a/pkg/runner/app.go b/pkg/runner/app.go index 0e340944..8b7ba75f 100644 --- a/pkg/runner/app.go +++ b/pkg/runner/app.go @@ -11,6 +11,8 @@ import ( "github.com/chainreactors/aiscan/agent" "github.com/chainreactors/aiscan/agent/hooks" "github.com/chainreactors/aiscan/agent/probe" + aop "github.com/chainreactors/aiscan/aop" + toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/capability" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" @@ -35,8 +37,12 @@ type App struct { SkillDiagnostics []skills.Diagnostic IOAClient *ioaclient.Client IOAStreamClient ioaclient.StreamAPI - DataBus *eventbus.Bus[output.ToolDataEvent] - Artifacts *output.ArtifactStream + EventBus *eventbus.Bus[*aop.Event] + Events *sessionEmitter + Progress *eventbus.Bus[*toolpb.Progress] + Recorder *output.JSONLRecorder + recorderMu sync.Mutex + closeOnce sync.Once enginesReady chan struct{} loggerMu sync.RWMutex logger telemetry.Logger @@ -55,8 +61,9 @@ func NewApp(ctx context.Context, rc ApplicationConfig) (*App, error) { a.Logger().Warnf("hook failed kind=%s source=%s error=%q", he.Kind, he.Source, he.Err) }) - a.DataBus = eventbus.New[output.ToolDataEvent]() - a.Artifacts = output.NewArtifactStream(a.DataBus) + a.EventBus = eventbus.New[*aop.Event]() + a.Events = newSessionEmitter(a.EventBus) + a.Progress = eventbus.New[*toolpb.Progress]() store, diagnostics := skills.LoadAll(rc.CLISkillPaths) a.Skills = store @@ -88,7 +95,13 @@ func NewApp(ctx context.Context, rc ApplicationConfig) (*App, error) { } } - a.Commands = initCoreCommands(rc, a.Provider, a.Skills, a.Hooks, logger) + a.Commands = initCoreCommands(rc, a.Provider, a.Skills, a.Hooks, a.Events, logger) + if rc.RecordFile != "" { + if err := a.StartRecording(rc.RecordFile); err != nil { + a.Close() + return nil, err + } + } a.enginesReady = make(chan struct{}) go func() { @@ -168,24 +181,71 @@ func (a *App) Close() { if a == nil { return } - if a.Artifacts != nil { - a.Artifacts.Close() - } - if a.Commands != nil { - for _, t := range a.Commands.Tools() { - if closer, ok := t.(interface{ Close() }); ok { - closer.Close() + a.closeOnce.Do(func() { + a.recorderMu.Lock() + if a.Recorder != nil { + if err := a.Recorder.Close(); err != nil { + a.Logger().Warnf("close AOP JSONL recorder: %s", err) } + a.Recorder = nil } - for _, cmd := range a.Commands.All() { - if cmd.Close != nil { - cmd.Close() + a.recorderMu.Unlock() + if a.Commands != nil { + for _, t := range a.Commands.Tools() { + if closer, ok := t.(interface{ Close() }); ok { + closer.Close() + } + } + for _, cmd := range a.Commands.All() { + if cmd.Close != nil { + cmd.Close() + } } } + if closer, ok := a.Engines.(interface{ Close() }); ok { + closer.Close() + } + }) +} + +func (a *App) StartRecording(path string) error { + if a == nil || strings.TrimSpace(path) == "" { + return nil } - if closer, ok := a.Engines.(interface{ Close() }); ok { - closer.Close() + a.recorderMu.Lock() + defer a.recorderMu.Unlock() + if a.Recorder != nil { + if !samePath(a.Recorder.Path(), path) { + return fmt.Errorf("AOP JSONL already records to %s", a.Recorder.Path()) + } + return nil + } + recorder, err := output.NewJSONLRecorder(a.EventBus, path) + if err != nil { + return err + } + a.Recorder = recorder + return nil +} + +func (a *App) SwitchRecording(path string) error { + if a == nil || strings.TrimSpace(path) == "" { + return fmt.Errorf("AOP JSONL path is required") + } + a.recorderMu.Lock() + defer a.recorderMu.Unlock() + if a.Recorder == nil { + recorder, err := output.NewJSONLRecorder(a.EventBus, path) + if err != nil { + return err + } + a.Recorder = recorder + return nil + } + if samePath(a.Recorder.Path(), path) { + return nil } + return a.Recorder.Switch(path) } func initProvider(provCfg agent.ProviderConfig, logger telemetry.Logger) (agent.Provider, *agent.ProviderConfig, error) { @@ -241,7 +301,7 @@ func llmConfigLabel(providerName, model string) string { return providerName + "/" + model } -func initCoreCommands(rc ApplicationConfig, llmProvider agent.Provider, skillStore *skills.Store, hookRegistry *hooks.Registry, logger telemetry.Logger) *commands.CommandRegistry { +func initCoreCommands(rc ApplicationConfig, llmProvider agent.Provider, skillStore *skills.Store, hookRegistry *hooks.Registry, events aop.EventEmitter, logger telemetry.Logger) *commands.CommandRegistry { cmdReg := commands.NewRegistry() workDir, _ := os.Getwd() deps := &commands.Deps{ @@ -253,6 +313,7 @@ func initCoreCommands(rc ApplicationConfig, llmProvider agent.Provider, skillSto TavilyKeys: rc.Tools.TavilyKeys, PlaywrightSession: rc.Tools.PlaywrightSession, Hooks: hookRegistry, + Events: events, } plan := capability.Select(capability.Options{ Groups: []string{"core", "arsenal", "search", "browser"}, diff --git a/pkg/runner/app_test.go b/pkg/runner/app_test.go index daa7335b..d6d7b1c7 100644 --- a/pkg/runner/app_test.go +++ b/pkg/runner/app_test.go @@ -1,16 +1,24 @@ package runner import ( + "bufio" "bytes" "context" "encoding/json" "net/http" "net/http/httptest" + "os" + "path/filepath" "strings" "testing" "github.com/chainreactors/aiscan/agent" + aop "github.com/chainreactors/aiscan/aop" + toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/telemetry" + "github.com/chainreactors/utils/parsers" + "google.golang.org/protobuf/encoding/protojson" + "google.golang.org/protobuf/types/known/anypb" ) func TestLogLLMProbeStatusReady(t *testing.T) { @@ -91,3 +99,82 @@ func TestAppLoggerCanBeRetargeted(t *testing.T) { t.Fatalf("retargeted logger missing: %q", second.String()) } } + +func TestJSONLRecorderPersistsCanonicalEventsAndOneArtifactPerResult(t *testing.T) { + path := filepath.Join(t.TempDir(), "session.jsonl") + app, err := NewApp(context.Background(), ApplicationConfig{ + RecordFile: path, SkipEngines: true, Logger: telemetry.NopLogger(), + }) + if err != nil { + t.Fatalf("NewApp: %v", err) + } + + app.Events.Emit(&aop.Event{ + SessionId: "session-1", TurnId: "turn-1", Emitter: "aiscan", + Payload: &aop.Event_ToolCall{ToolCall: &aop.ToolCall{Id: "call-1", Name: "gogo"}}, + }) + app.Progress.Emit(&toolpb.Progress{Tool: "gogo", Text: "raw PTY bytes", CallId: "call-1"}) + gogoResult := parsers.NewGOGOResult("127.0.0.1", "443") + gogoResult.Protocol = "https" + raw, err := json.Marshal(gogoResult) + if err != nil { + t.Fatal(err) + } + extension, err := anypb.New(&toolpb.Artifact{ + Tool: "gogo", Kind: toolpb.ArtifactKindService, Target: gogoResult.GetTarget(), Data: raw, + MediaType: aop.JSONMediaType, CallId: "call-1", + }) + if err != nil { + t.Fatal(err) + } + app.Events.Emit(&aop.Event{ + SessionId: "session-1", TurnId: "turn-1", Emitter: "aiscan", + Payload: &aop.Event_Extension{Extension: extension}, + }) + app.Events.Emit(&aop.Event{ + SessionId: "session-1", TurnId: "turn-1", Emitter: "aiscan", + Payload: &aop.Event_ToolResult{ToolResult: &aop.ToolResult{CallId: "call-1", Name: "gogo"}}, + }) + app.Close() + + file, err := os.Open(path) + if err != nil { + t.Fatalf("open JSONL: %v", err) + } + defer file.Close() + counts := map[string]int{} + var artifact toolpb.Artifact + scanner := bufio.NewScanner(file) + for scanner.Scan() { + line := scanner.Bytes() + if len(bytes.TrimSpace(line)) == 0 { + t.Fatal("JSONL contains a blank line") + } + event := new(aop.Event) + if err := protojson.Unmarshal(line, event); err != nil { + t.Fatalf("JSONL line is not an AOP event: %s", err) + } + counts[aop.Kind(event)]++ + if extension := event.GetExtension(); extension != nil && extension.MessageIs(&artifact) { + if err := extension.UnmarshalTo(&artifact); err != nil { + t.Fatalf("decode artifact: %v", err) + } + } + } + if err := scanner.Err(); err != nil { + t.Fatalf("read JSONL: %v", err) + } + if counts["tool.call"] != 1 || counts["tool.result"] != 1 || counts["aop.tool.Artifact"] != 1 { + t.Fatalf("event counts = %#v", counts) + } + if artifact.Tool != "gogo" || artifact.Kind != toolpb.ArtifactKindService || artifact.Target != "127.0.0.1:443" || artifact.CallId != "call-1" { + t.Fatalf("artifact = %#v", &artifact) + } + var decoded parsers.GOGOResult + if err := json.Unmarshal(artifact.Data, &decoded); err != nil { + t.Fatalf("decode gogo result: %v", err) + } + if decoded.Ip != "127.0.0.1" || decoded.Port != "443" || decoded.Protocol != "https" { + t.Fatalf("gogo result = %#v", decoded) + } +} diff --git a/pkg/runner/application_builder.go b/pkg/runner/application_builder.go index d751d1ce..d3b31e4a 100644 --- a/pkg/runner/application_builder.go +++ b/pkg/runner/application_builder.go @@ -53,9 +53,8 @@ func AppConfigFromDistribute(dc *types.DistributeConfig, features RuntimeFeature } } -// MergeOptionExtras layers the fields DistributeConfig does not model onto a -// proto-built ApplicationConfig: playwright session, uncover credentials, and -// CLI skill paths. +// MergeOptionExtras layers fields DistributeConfig does not model onto a +// proto-built ApplicationConfig. func MergeOptionExtras(rc ApplicationConfig, option *cfg.Option) ApplicationConfig { if option == nil { return rc @@ -63,6 +62,7 @@ func MergeOptionExtras(rc ApplicationConfig, option *cfg.Option) ApplicationConf rc.Scanner.UncoverCredentials = cloneStringMap(option.UncoverCredentials) rc.Tools.PlaywrightSession = option.PlaywrightSession rc.CLISkillPaths = skillPathsFromOptions(option) + rc.RecordFile = option.OutputFile return rc } @@ -98,6 +98,7 @@ func AppConfig(option *cfg.Option, features RuntimeFeatures, logger telemetry.Lo }, Logger: logger, CLISkillPaths: skillPathsFromOptions(option), + RecordFile: option.OutputFile, } } diff --git a/pkg/runner/application_config.go b/pkg/runner/application_config.go index 096eb196..ff324d80 100644 --- a/pkg/runner/application_config.go +++ b/pkg/runner/application_config.go @@ -13,6 +13,7 @@ type ApplicationConfig struct { IOA *IOAConfig Logger telemetry.Logger CLISkillPaths []string + RecordFile string SkipEngines bool } diff --git a/pkg/runner/local_repl.go b/pkg/runner/local_repl.go index a27a8f37..9dbc6ca9 100644 --- a/pkg/runner/local_repl.go +++ b/pkg/runner/local_repl.go @@ -31,7 +31,7 @@ func (rt *AgentRuntime) AttachLocalREPL(ctx context.Context) error { ctx, option, rt.consoleAppInfoForSession(sess), - sess.state.agent, + sess.Agent(), rlterm.Local(), rt.Subscribe, ) diff --git a/pkg/runner/provider_config_test.go b/pkg/runner/provider_config_test.go index 9c8a1d02..2c3f0aca 100644 --- a/pkg/runner/provider_config_test.go +++ b/pkg/runner/provider_config_test.go @@ -109,9 +109,10 @@ func TestMergeOptionExtrasLayersNonProtoFields(t *testing.T) { option := &cfg.Option{ PlaywrightSession: "browser-1", UncoverCredentials: map[string]string{"SHODAN_API_KEY": "shodan-key"}, + MiscOptions: cfg.MiscOptions{OutputFile: "session.jsonl"}, } rc = MergeOptionExtras(rc, option) - if rc.Tools.PlaywrightSession != "browser-1" || rc.Scanner.UncoverCredentials["SHODAN_API_KEY"] != "shodan-key" { + if rc.Tools.PlaywrightSession != "browser-1" || rc.Scanner.UncoverCredentials["SHODAN_API_KEY"] != "shodan-key" || rc.RecordFile != "session.jsonl" { t.Fatalf("extras = %+v", rc) } } diff --git a/pkg/runner/remote_repl.go b/pkg/runner/remote_repl.go index fd3f640d..07d04d45 100644 --- a/pkg/runner/remote_repl.go +++ b/pkg/runner/remote_repl.go @@ -36,7 +36,7 @@ func (rt *AgentRuntime) startMainREPL() error { Resize: control.SetSize, }, func(replCtx context.Context, input io.Reader, output io.Writer) error { for { - err := tui.RunRemoteAgentConsoleWithControl(replCtx, option, rt.consoleAppInfoForSession(sess), sess.state.agent, input, output, control, rt.Subscribe) + err := tui.RunRemoteAgentConsoleWithControl(replCtx, option, rt.consoleAppInfoForSession(sess), sess.Agent(), input, output, control, rt.Subscribe) if replCtx.Err() != nil { return replCtx.Err() } diff --git a/pkg/runner/runner.go b/pkg/runner/runner.go index a948e3ac..5dc8d69d 100644 --- a/pkg/runner/runner.go +++ b/pkg/runner/runner.go @@ -3,8 +3,11 @@ package runner import ( "context" "encoding/json" + "errors" "fmt" "os" + "path/filepath" + "runtime" "strings" "sync" "time" @@ -30,31 +33,33 @@ import ( // --------------------------------------------------------------------------- type AgentRuntime struct { - app *App - nodeName string - systemPrompt string - option *cfg.Option - config agent.Config - bus *eventbus.Bus[*aop.Event] - sessionEvents *sessionEmitter - output RunOutput - configFile string - resumeMessages []*aop.Message - ctx context.Context - cancel context.CancelFunc - mu sync.RWMutex - sessions map[string]*sessionState - runs map[string]*Run - requestSeq uint64 - closeOnce sync.Once - wg sync.WaitGroup - operations sync.WaitGroup - namespaceMux *aop.NamespaceMux - ptyManager *tmuxpkg.Manager - replMode REPLMode - maxPending int - ownsApp bool - cleanup func() + app *App + nodeName string + systemPrompt string + option *cfg.Option + config agent.Config + bus *eventbus.Bus[*aop.Event] + sessionEvents *sessionEmitter + output RunOutput + configFile string + resumeMessages []*aop.Message + resumeSessionID string + recordPath string + ctx context.Context + cancel context.CancelFunc + mu sync.RWMutex + sessions map[string]*sessionState + runs map[string]*Run + requestSeq uint64 + closeOnce sync.Once + wg sync.WaitGroup + operations sync.WaitGroup + namespaceMux *aop.NamespaceMux + ptyManager *tmuxpkg.Manager + replMode REPLMode + maxPending int + ownsApp bool + cleanup func() } type REPLMode uint8 @@ -65,6 +70,37 @@ const ( REPLPersistent ) +func resolveJSONLRecordPath(option *cfg.Option, replMode REPLMode) string { + if option == nil { + return "" + } + if path := strings.TrimSpace(option.Resume); path != "" { + return path + } + if path := strings.TrimSpace(option.OutputFile); path != "" { + return path + } + if !option.SaveSession && replMode == REPLDisabled { + return "" + } + name := "session-" + time.Now().Format("20060102-150405.000000000") + ".jsonl" + return filepath.Join(cfg.DataSubDir("sessions"), name) +} + +func samePath(left, right string) bool { + leftAbs, leftErr := filepath.Abs(left) + rightAbs, rightErr := filepath.Abs(right) + if leftErr == nil && rightErr == nil { + left, right = filepath.Clean(leftAbs), filepath.Clean(rightAbs) + } else { + left, right = filepath.Clean(left), filepath.Clean(right) + } + if runtime.GOOS == "windows" { + return strings.EqualFold(left, right) + } + return left == right +} + // 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. @@ -117,6 +153,12 @@ func NewAgentRuntime(ctx context.Context, option *cfg.Option, logger telemetry.L rt.option = &optCopy rt.configFile = option.ConfigFile } + recordPath := resolveJSONLRecordPath(option, rt.replMode) + if recordPath != "" { + option.OutputFile = recordPath + rt.option.OutputFile = recordPath + rt.recordPath = recordPath + } if rc != nil && rc.ExistingApp != nil { rt.app = rc.ExistingApp @@ -161,6 +203,35 @@ func NewAgentRuntime(ctx context.Context, option *cfg.Option, logger telemetry.L rt.app.SetLogger(logger) logger = rt.app.Logger() } + publicBus := rt.app.EventBus + if publicBus == nil { + publicBus = eventbus.New[*aop.Event]() + rt.app.EventBus = publicBus + } + rt.sessionEvents = rt.app.Events + if rt.sessionEvents == nil { + rt.sessionEvents = newSessionEmitter(publicBus) + rt.app.Events = rt.sessionEvents + } + if recordPath != "" { + if err := rt.app.StartRecording(recordPath); err != nil { + rt.Close() + return nil, fmt.Errorf("open JSONL recorder: %w", err) + } + logger.Importantf("recording session JSONL to %s", recordPath) + } + var resumeCounter int64 + if option.Resume != "" { + data, err := loadResumeState(option.Resume) + if err != nil { + rt.Close() + return nil, fmt.Errorf("resume session: %w", err) + } + rt.resumeMessages = data.Messages + rt.resumeSessionID = data.SessionID + resumeCounter = data.MessageCounter + logger.Importantf("resumed %d messages from %s", len(data.Messages), option.Resume) + } nodeName := ResolveIOANodeName(option) rt.nodeName = nodeName @@ -203,12 +274,10 @@ func NewAgentRuntime(ctx context.Context, option *cfg.Option, logger telemetry.L rt.output = rc.Output } - publicBus := eventbus.New[*aop.Event]() if rt.output != nil { publicBus.Subscribe(rt.output.HandleEvent) } rt.bus = publicBus - rt.sessionEvents = newSessionEmitter(publicBus) var ioaCancel func() var handoffCancel func() @@ -238,23 +307,7 @@ func NewAgentRuntime(ctx context.Context, option *cfg.Option, logger telemetry.L Bus: rt.sessionEvents, Hooks: rt.app.Hooks, CaptureProviderFrames: option.CaptureProviderFrames, - } - - if option.SaveSession { - sessDir := cfg.DataSubDir("sessions") - rt.config = rt.config.WithOnRunEnd(func(result *agent.Result) { - if result == nil || len(result.Messages) == 0 { - return - } - if err := agent.SaveCheckpoint(sessDir, &agent.CheckpointData{ - Model: option.Model, - Provider: option.Provider, - Messages: result.Messages, - MessageCounter: result.MessageCounter, - }); err != nil { - logger.Warnf("save session: %s", err) - } - }) + MessageCounter: resumeCounter, } subAgentTool := agent.NewSubAgentTool(func(name string) (agent.AgentType, error) { @@ -287,17 +340,6 @@ func NewAgentRuntime(ctx context.Context, option *cfg.Option, logger telemetry.L Run: loop.Run, }, "loop") - if option.Resume != "" { - path := option.Resume - data, err := agent.LoadCheckpoint(path) - if err != nil { - rt.Close() - return nil, fmt.Errorf("resume session: %w", err) - } - rt.resumeMessages = data.Messages - logger.Importantf("resumed %d messages from %s", len(data.Messages), path) - } - if rt.app.IOAStreamClient != nil && option.Space != "" { nodeID := "" if rt.app.IOAClient != nil { @@ -526,7 +568,7 @@ func runInteractiveMode(ctx context.Context, option *cfg.Option, logger telemetr // Scanner direct execution // --------------------------------------------------------------------------- -func RunDirectScannerMode(ctx context.Context, option *cfg.Option, rest []string, logger telemetry.Logger) error { +func RunDirectScannerMode(ctx context.Context, option *cfg.Option, rest []string, logger telemetry.Logger) (runErr error) { defaultVerify := cfg.ResolveString(option.ScanConfig.Verify, cfg.DefaultVerify) features, scannerArgs, err := DirectScannerRuntimeFeaturesWithDefault(rest, defaultVerify) if err != nil { @@ -550,6 +592,9 @@ func RunDirectScannerMode(ctx context.Context, option *cfg.Option, rest []string return nil } } + if recordPath := resolveJSONLRecordPath(option, REPLDisabled); recordPath != "" { + option.OutputFile = recordPath + } scannerLogger := logger if !directScannerDebugEnabled(option, scannerArgs) { @@ -585,6 +630,9 @@ func RunDirectScannerMode(ctx context.Context, option *cfg.Option, rest []string if option.NoColor && scannerArgs[0] == "scan" && !HasScannerFlag(scannerArgs[1:], "--no-color") { scannerArgs = append(scannerArgs, "--no-color") } + sessionID := fmt.Sprintf("scan-%d", time.Now().UnixNano()) + turnID := sessionID + "-run" + emitter := scannerArgs[0] tool, ok := application.Commands.GetTool("bash") if !ok { return fmt.Errorf("bash tool is not registered") @@ -593,6 +641,53 @@ func RunDirectScannerMode(ctx context.Context, option *cfg.Option, rest []string if !ok { return fmt.Errorf("registered bash tool has unexpected type") } + callID := turnID + "-call" + ctx = coretool.ContextWithInvocation(ctx, coretool.Invocation{ + CallID: callID, SessionID: sessionID, TurnID: turnID, Emitter: emitter, + }) + arguments, err := aop.JSONValue(map[string]any{"args": scannerArgs[1:]}) + if err != nil { + return fmt.Errorf("encode scanner arguments: %w", err) + } + startedAt := time.Now() + if application.Events != nil { + application.Events.sessionStarted(sessionID, emitter, &aop.SessionStarted{}) + application.Events.Emit(&aop.Event{ + SessionId: sessionID, TurnId: turnID, Emitter: emitter, + Payload: &aop.Event_TurnStarted{TurnStarted: &aop.TurnStarted{}}, + }) + application.Events.Emit(&aop.Event{ + SessionId: sessionID, TurnId: turnID, Emitter: emitter, + Payload: &aop.Event_ToolCall{ToolCall: &aop.ToolCall{Id: callID, Name: emitter, Arguments: arguments}}, + }) + defer func() { + isCanceled := errors.Is(runErr, context.Canceled) || errors.Is(ctx.Err(), context.Canceled) + result := &aop.ToolResult{ + CallId: callID, Name: emitter, IsError: runErr != nil, + DurationMs: uint64(time.Since(startedAt).Milliseconds()), + } + stopReason := string(agent.StopReasonCompleted) + closeReason := SessionCloseCompleted + if runErr != nil { + result.Output = []*aop.Content{aop.Text(runErr.Error())} + stopReason = string(agent.StopReasonError) + closeReason = SessionCloseError + } + if isCanceled { + stopReason = string(agent.StopReasonCanceled) + closeReason = SessionCloseCanceled + } + application.Events.Emit(&aop.Event{ + SessionId: sessionID, TurnId: turnID, Emitter: emitter, + Payload: &aop.Event_ToolResult{ToolResult: result}, + }) + application.Events.Emit(&aop.Event{ + SessionId: sessionID, TurnId: turnID, Emitter: emitter, + Payload: &aop.Event_TurnEnded{TurnEnded: &aop.TurnEnded{StopReason: stopReason}}, + }) + application.Events.sessionEnded(sessionID, emitter, string(closeReason)) + }() + } streaming := ShouldStreamScannerOutput(scannerArgs) var captured strings.Builder execution, err := bash.RunForeground(ctx, cmdpkg.JoinCommandLine(scannerArgs[0], scannerArgs[1:]), cmdpkg.BashExecOptions{ @@ -607,6 +702,9 @@ func RunDirectScannerMode(ctx context.Context, option *cfg.Option, rest []string if err != nil { return err } + if ctx.Err() != nil { + return ctx.Err() + } if !streaming { fmt.Print(captured.String()) } diff --git a/pkg/runner/runner_test.go b/pkg/runner/runner_test.go new file mode 100644 index 00000000..af5e6450 --- /dev/null +++ b/pkg/runner/runner_test.go @@ -0,0 +1,401 @@ +package runner + +import ( + "context" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/agent/provider" + aop "github.com/chainreactors/aiscan/aop" + cfg "github.com/chainreactors/aiscan/core/config" + "github.com/chainreactors/aiscan/core/eventbus" + "github.com/chainreactors/aiscan/core/output" + "github.com/chainreactors/aiscan/core/telemetry" + "google.golang.org/protobuf/types/known/timestamppb" +) + +type persistenceProvider struct { + requests []*provider.ChatCompletionRequest +} + +func (p *persistenceProvider) Name() string { return "persistence" } + +func (p *persistenceProvider) ChatCompletion(_ context.Context, request *provider.ChatCompletionRequest) (*provider.ChatCompletionResponse, error) { + p.requests = append(p.requests, request) + return &provider.ChatCompletionResponse{ + Choices: []provider.Choice{{Message: provider.TextMessage("assistant", "persisted response")}}, + }, nil +} + +func TestFileFlagPersistsOneCanonicalAOPStream(t *testing.T) { + path := filepath.Join(t.TempDir(), "explicit.jsonl") + option := &cfg.Option{MiscOptions: cfg.MiscOptions{OutputFile: path}} + provider := new(persistenceProvider) + app, runtime := newPersistenceRuntime(t, option, provider) + + session, err := runtime.OpenSession(context.Background(), SessionOptions{ID: "task"}) + if err != nil { + t.Fatal(err) + } + run, err := session.Run(context.Background(), RunInput{Content: []*aop.Content{aop.Text("persist this")}}) + if err != nil { + t.Fatal(err) + } + if _, err := run.Wait(); err != nil { + t.Fatal(err) + } + if err := runtime.CloseSession(context.Background(), "task", SessionCloseCompleted); err != nil { + t.Fatal(err) + } + runtime.Close() + app.Close() + + events, err := output.ReadJSONL(path) + if err != nil { + t.Fatalf("ReadJSONL: %v", err) + } + counts := map[string]int{} + for _, event := range events { + if event.SessionId == "" || event.Payload == nil { + t.Fatalf("invalid AOP event: %#v", event) + } + counts[aop.Kind(event)]++ + } + for _, kind := range []string{"session.started", "turn.started", "message", "turn.ended", "session.ended"} { + if counts[kind] == 0 { + t.Fatalf("missing %s in %#v", kind, counts) + } + } + if counts["message"] != 2 { + t.Fatalf("message count = %d, want user + assistant", counts["message"]) + } +} + +func TestResumeRestoresAndAppendsAOPStream(t *testing.T) { + dir := t.TempDir() + resumePath := filepath.Join(dir, "resume.jsonl") + writePersistenceSession(t, resumePath) + baseEvents, err := output.ReadJSONL(resumePath) + if err != nil { + t.Fatal(err) + } + + t.Run("append same file", func(t *testing.T) { + option := &cfg.Option{} + option.Resume = resumePath + provider := new(persistenceProvider) + app, runtime := newPersistenceRuntime(t, option, provider) + runResumedTurn(t, runtime, "continued prompt") + runtime.Close() + app.Close() + + if len(provider.requests) != 1 { + t.Fatalf("provider requests = %d", len(provider.requests)) + } + requestText := persistenceRequestText(provider.requests[0]) + for _, expected := range []string{"old user", "old assistant", "continued prompt"} { + if !strings.Contains(requestText, expected) { + t.Fatalf("resumed request missing %q:\n%s", expected, requestText) + } + } + events, err := output.ReadJSONL(resumePath) + if err != nil { + t.Fatal(err) + } + if len(events) <= len(baseEvents) { + t.Fatalf("resume did not append: before=%d after=%d", len(baseEvents), len(events)) + } + data, err := loadResumeState(resumePath) + if err != nil { + t.Fatal(err) + } + if data.MessageCounter < 4 || len(data.Messages) != 4 { + t.Fatalf("resumed session = %#v", data) + } + }) +} + +func TestREPLResumeLoadsMainSessionContext(t *testing.T) { + resumePath := filepath.Join(t.TempDir(), "repl-resume.jsonl") + writePersistenceSessionForID(t, resumePath, MainREPLName) + option := &cfg.Option{} + option.Resume = resumePath + provider := new(persistenceProvider) + app, runtime := newPersistenceRuntimeWithMode(t, option, provider, REPLEphemeral) + + session, err := runtime.OpenSession(context.Background(), SessionOptions{ID: MainREPLName}) + if err != nil { + t.Fatal(err) + } + messages := session.MessagesSnapshot() + if len(messages) != 2 { + t.Fatalf("REPL resumed messages = %d, want 2", len(messages)) + } + text := persistenceMessagesText(messages) + if !strings.Contains(text, "old user") || !strings.Contains(text, "old assistant") { + t.Fatalf("REPL context = %q", text) + } + if err := runtime.CloseSession(context.Background(), MainREPLName, SessionCloseCompleted); err != nil { + t.Fatal(err) + } + runtime.Close() + app.Close() +} + +func TestClearRotatesToAnEmptyContinuationSession(t *testing.T) { + path := filepath.Join(t.TempDir(), "clear.jsonl") + option := &cfg.Option{MiscOptions: cfg.MiscOptions{OutputFile: path}} + provider := new(persistenceProvider) + app, runtime := newPersistenceRuntimeWithMode(t, option, provider, REPLEphemeral) + session, err := runtime.OpenSession(context.Background(), SessionOptions{ID: MainREPLName}) + if err != nil { + t.Fatal(err) + } + run, err := session.Run(context.Background(), RunInput{Content: []*aop.Content{aop.Text("before clear")}}) + if err != nil { + t.Fatal(err) + } + if _, err := run.Wait(); err != nil { + t.Fatal(err) + } + oldID := session.ID() + + var events []*aop.Event + unsub := runtime.Subscribe(func(event *aop.Event) { events = append(events, event) }) + result, err := session.Command(context.Background(), "/clear") + unsub() + if err != nil { + t.Fatalf("/clear: %v", err) + } + if text := persistenceMessagesText([]*aop.Message{{Content: result.Content}}); !strings.Contains(text, "Context cleared") { + t.Fatalf("clear result = %#v", result) + } + newID := session.ID() + if newID == "" || newID == oldID { + t.Fatalf("clear session id = %q, old = %q", newID, oldID) + } + if messages := session.MessagesSnapshot(); len(messages) != 0 { + t.Fatalf("new clear context has %d messages", len(messages)) + } + assertRotationEvents(t, events, oldID, newID, string(SessionCloseCleared)) + + if err := runtime.CloseSession(context.Background(), MainREPLName, SessionCloseCompleted); err != nil { + t.Fatal(err) + } + runtime.Close() + app.Close() + data, err := loadResumeState(path) + if err != nil { + t.Fatalf("LoadSession after clear: %v", err) + } + if data.SessionID != newID || len(data.Messages) != 0 { + t.Fatalf("clear resume data = %#v", data) + } +} + +func TestCompactRotatesAndPersistsOnlyCompactedContext(t *testing.T) { + path := filepath.Join(t.TempDir(), "compact.jsonl") + option := &cfg.Option{MiscOptions: cfg.MiscOptions{OutputFile: path}} + provider := new(persistenceProvider) + app, runtime := newPersistenceRuntimeWithMode(t, option, provider, REPLEphemeral) + runtime.config.Compaction = agent.CompactionSettings{KeepRecentTokens: 20, ReserveTokens: 64} + long := strings.Repeat("history ", 120) + messages := []*aop.Message{ + agent.TextMessage("user", long+"one"), + agent.TextMessage("assistant", long+"two"), + agent.TextMessage("user", long+"three"), + agent.TextMessage("assistant", "recent answer"), + } + session, err := runtime.OpenSession(context.Background(), SessionOptions{ID: MainREPLName, Messages: messages}) + if err != nil { + t.Fatal(err) + } + oldID := session.ID() + var events []*aop.Event + unsub := runtime.Subscribe(func(event *aop.Event) { events = append(events, event) }) + if _, err := session.Command(context.Background(), "/compact focus on findings"); err != nil { + unsub() + t.Fatalf("/compact: %v", err) + } + unsub() + newID := session.ID() + if newID == oldID { + t.Fatal("compact did not rotate the session") + } + compacted := session.MessagesSnapshot() + if len(compacted) == 0 || len(compacted) >= len(messages) { + t.Fatalf("compacted messages = %d, original = %d", len(compacted), len(messages)) + } + assertRotationEvents(t, events, oldID, newID, string(SessionCloseCompacted)) + + if err := runtime.CloseSession(context.Background(), MainREPLName, SessionCloseCompleted); err != nil { + t.Fatal(err) + } + runtime.Close() + app.Close() + data, err := loadResumeState(path) + if err != nil { + t.Fatal(err) + } + if data.SessionID != newID || len(data.Messages) != len(compacted) { + t.Fatalf("compacted resume data = %#v, want %d messages", data, len(compacted)) + } + if strings.Contains(persistenceMessagesText(data.Messages), long+"one") { + t.Fatal("compacted resume context retained discarded history") + } +} + +func TestInteractiveResumeRotatesAndBootstrapsSelectedContext(t *testing.T) { + dir := t.TempDir() + currentPath := filepath.Join(dir, "current.jsonl") + resumePath := filepath.Join(dir, "selected.jsonl") + writePersistenceSessionForID(t, resumePath, "selected-main") + option := &cfg.Option{MiscOptions: cfg.MiscOptions{OutputFile: currentPath}} + provider := new(persistenceProvider) + app, runtime := newPersistenceRuntimeWithMode(t, option, provider, REPLEphemeral) + session, err := runtime.OpenSession(context.Background(), SessionOptions{ID: MainREPLName}) + if err != nil { + t.Fatal(err) + } + oldID := session.ID() + count, err := session.Resume(context.Background(), resumePath) + if err != nil { + t.Fatalf("Resume: %v", err) + } + if count != 2 { + t.Fatalf("resumed messages = %d", count) + } + newID := session.ID() + if newID == oldID { + t.Fatal("interactive resume did not rotate the session") + } + if text := persistenceMessagesText(session.MessagesSnapshot()); !strings.Contains(text, "old user") || !strings.Contains(text, "old assistant") { + t.Fatalf("resumed context = %q", text) + } + state := session.currentState() + if state.parentSessionID != "selected-main" || state.parentToolCallID != "" { + t.Fatalf("resumed continuation parent = %q/%q", state.parentSessionID, state.parentToolCallID) + } + + if err := runtime.CloseSession(context.Background(), MainREPLName, SessionCloseCompleted); err != nil { + t.Fatal(err) + } + runtime.Close() + app.Close() + data, err := loadResumeState(resumePath) + if err != nil { + t.Fatal(err) + } + if data.SessionID != newID || len(data.Messages) != 2 { + t.Fatalf("interactive resume JSONL = %#v", data) + } + currentEvents, err := output.ReadJSONL(currentPath) + if err != nil { + t.Fatal(err) + } + if len(currentEvents) == 0 || currentEvents[len(currentEvents)-1].GetSessionEnded().GetReason() != string(SessionCloseResumed) { + t.Fatalf("current file was not closed before switch: %#v", currentEvents) + } +} + +func assertRotationEvents(t *testing.T, events []*aop.Event, oldID, newID, reason string) { + t.Helper() + var ended, started bool + for _, event := range events { + if event.SessionId == oldID && event.GetSessionEnded().GetReason() == reason { + ended = true + } + if event.SessionId == newID && event.GetSessionStarted().GetParentSessionId() == oldID && event.GetSessionStarted().GetParentToolCallId() == "" { + started = true + } + } + if !ended || !started { + t.Fatalf("rotation events ended=%v started=%v events=%#v", ended, started, events) + } +} + +func newPersistenceRuntime(t *testing.T, option *cfg.Option, llm *persistenceProvider) (*App, *AgentRuntime) { + return newPersistenceRuntimeWithMode(t, option, llm, REPLDisabled) +} + +func newPersistenceRuntimeWithMode(t *testing.T, option *cfg.Option, llm *persistenceProvider, replMode REPLMode) (*App, *AgentRuntime) { + t.Helper() + app, err := NewApp(context.Background(), ApplicationConfig{SkipEngines: true, Logger: telemetry.NopLogger()}) + if err != nil { + t.Fatal(err) + } + app.Provider = llm + app.ProviderConfig = agent.ProviderConfig{Provider: llm.Name(), Model: "test-model", MaxTokens: 128, ContextWindow: 128000} + runtime, err := NewAgentRuntime(context.Background(), option, telemetry.NopLogger(), &RuntimeConfig{ExistingApp: app, REPLMode: replMode}) + if err != nil { + app.Close() + t.Fatal(err) + } + t.Cleanup(func() { + runtime.Close() + app.Close() + }) + return app, runtime +} + +func runResumedTurn(t *testing.T, runtime *AgentRuntime, prompt string) { + t.Helper() + session, err := runtime.OpenSession(context.Background(), SessionOptions{ID: "task", Messages: runtime.resumeMessages}) + if err != nil { + t.Fatal(err) + } + run, err := session.Run(context.Background(), RunInput{Content: []*aop.Content{aop.Text(prompt)}}) + if err != nil { + t.Fatal(err) + } + if _, err := run.Wait(); err != nil { + t.Fatal(err) + } + if err := runtime.CloseSession(context.Background(), "task", SessionCloseCompleted); err != nil { + t.Fatal(err) + } +} + +func writePersistenceSession(t *testing.T, path string) { + writePersistenceSessionForID(t, path, "task") +} + +func writePersistenceSessionForID(t *testing.T, path, sessionID string) { + t.Helper() + timestamp := timestamppb.New(time.Date(2026, 8, 3, 0, 0, 0, 0, time.UTC)) + events := []*aop.Event{ + {Id: "e-1", EmittedAt: timestamp, SessionId: sessionID, Emitter: "aiscan", Seq: 1, Payload: &aop.Event_SessionStarted{SessionStarted: &aop.SessionStarted{Model: "test-model"}}}, + {Id: "e-2", EmittedAt: timestamp, SessionId: sessionID, TurnId: "old-turn", Emitter: "aiscan", Seq: 2, Payload: &aop.Event_Message{Message: &aop.Message{Id: "m-1", Role: "user", Content: []*aop.Content{aop.Text("old user")}}}}, + {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"}}}, + } + bus := eventbus.New[*aop.Event]() + writer, err := output.NewJSONLRecorder(bus, path) + if err != nil { + t.Fatal(err) + } + for _, event := range events { + bus.Emit(event) + } + if err := writer.Close(); err != nil { + t.Fatal(err) + } +} + +func persistenceRequestText(request *provider.ChatCompletionRequest) string { + if request == nil { + return "" + } + return persistenceMessagesText(request.Messages) +} + +func persistenceMessagesText(messages []*aop.Message) string { + var parts []string + for _, message := range messages { + parts = append(parts, provider.MessageText(message)) + } + return strings.Join(parts, "\n") +} diff --git a/pkg/runner/runtime_session.go b/pkg/runner/runtime_session.go index 8df648d4..66391f1d 100644 --- a/pkg/runner/runtime_session.go +++ b/pkg/runner/runtime_session.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "fmt" + "strconv" "strings" "sync" "time" @@ -13,7 +14,9 @@ import ( "github.com/chainreactors/aiscan/agent/evaluator" inboxpkg "github.com/chainreactors/aiscan/agent/inbox" aop "github.com/chainreactors/aiscan/aop" + cfg "github.com/chainreactors/aiscan/core/config" "github.com/chainreactors/aiscan/core/eventbus" + "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/telemetry" toolpkg "github.com/chainreactors/aiscan/core/tool" "github.com/chainreactors/aiscan/pkg/commands" @@ -28,6 +31,7 @@ const DefaultSessionPendingLimit = 64 type SessionOptions struct { ID string + LogicalID string ParentSessionID string ParentToolCallID string AgentName string @@ -39,6 +43,10 @@ type SessionCloseReason string const ( SessionCloseCompleted SessionCloseReason = "completed" SessionCloseCanceled SessionCloseReason = "canceled" + SessionCloseError SessionCloseReason = "error" + SessionCloseCleared SessionCloseReason = "cleared" + SessionCloseCompacted SessionCloseReason = "compacted" + SessionCloseResumed SessionCloseReason = "resumed" SessionCloseRuntime SessionCloseReason = "runtime_closed" ) @@ -67,7 +75,9 @@ const ( ) type Session struct { - state *sessionState + mu sync.RWMutex + state *sessionState + logicalID string } type Run struct { @@ -128,7 +138,7 @@ func newSessionEmitter(bus *eventbus.Bus[*aop.Event]) *sessionEmitter { // Emit stamps the event with runtime metadata (timestamp, per-session // sequence, fallback id) and forwards it to the runtime's single public bus. -// It satisfies agent.EventEmitter so session agents emit through it directly. +// It satisfies aop.EventEmitter so session agents and tools emit through it directly. func (e *sessionEmitter) Emit(event *aop.Event) { if event.EmittedAt == nil { event.EmittedAt = timestamppb.Now() @@ -217,19 +227,6 @@ func (s *commandSession) execute(ctx context.Context, input string) commandOutco return commandText(line, CommandPresentationPreformatted, fmt.Sprintf( "Session: %s\nAgent: %s\nProvider: %s\nModel: %s\nMessages: %d", s.state.id, s.state.agentName, providerName, model, len(s.state.agent.MessagesSnapshot()))) - case "/clear": - s.state.agent.Reset() - return commandText(line, CommandPresentationPlain, "Context cleared.") - case "/compact": - if len(s.state.agent.MessagesSnapshot()) < 4 { - return commandText(line, CommandPresentationPlain, "Nothing to compact (too few messages).") - } - result, err := s.state.agent.Compact(ctx, agent.CompactConfig{CustomInstructions: strings.TrimSpace(strings.Join(values, " "))}) - if err != nil { - return commandOutcome{err: err} - } - return commandText(line, CommandPresentationPlain, fmt.Sprintf( - "Compacted: ~%d -> ~%d tokens (%d messages kept)", result.TokensBefore, result.TokensAfter, result.KeptMessages)) case "/eval", "/goal": criteria := strings.TrimSpace(strings.Join(values, " ")) switch criteria { @@ -348,6 +345,7 @@ func (m *sessionMailbox) ActiveProducers() int { return m.base.ActiveProducers() type sessionState struct { runtime *AgentRuntime id string + logicalID string agentName string parentSessionID string parentToolCallID string @@ -376,6 +374,19 @@ func (rt *AgentRuntime) OpenSession(ctx context.Context, options SessionOptions) if id == "" { id = rt.nextRuntimeID("session") } + logicalID := strings.TrimSpace(options.LogicalID) + if logicalID == "" { + logicalID = id + } + if rt.resumeSessionID != "" && options.LogicalID == "" && (logicalID == MainREPLName || logicalID == "task") { + id = rt.nextContinuationID(logicalID) + if options.ParentSessionID == "" { + options.ParentSessionID = rt.resumeSessionID + } + if len(options.Messages) == 0 { + options.Messages = rt.resumeMessages + } + } agentName := strings.TrimSpace(options.AgentName) if agentName == "" { agentName = rt.nodeName @@ -389,9 +400,9 @@ func (rt *AgentRuntime) OpenSession(ctx context.Context, options SessionOptions) rt.mu.Unlock() return nil, rt.ctx.Err() } - if _, exists := rt.sessions[id]; exists { + if _, exists := rt.sessions[logicalID]; exists { rt.mu.Unlock() - return nil, fmt.Errorf("session %q already exists", id) + return nil, fmt.Errorf("session %q already exists", logicalID) } sessionCtx, cancel := context.WithCancel(ctx) baseInbox := inboxpkg.NewBuffered(agent.DefaultInboxCapacity) @@ -414,20 +425,20 @@ func (rt *AgentRuntime) OpenSession(ctx context.Context, options SessionOptions) ag.LoadMessages(rt.resumeMessages) } state := &sessionState{ - runtime: rt, id: id, agentName: agentName, + runtime: rt, id: id, logicalID: logicalID, agentName: agentName, parentSessionID: options.ParentSessionID, parentToolCallID: options.ParentToolCallID, agent: ag, inbox: mailbox, scheduler: scheduler, ctx: sessionCtx, cancel: cancel, ops: make(chan *sessionOperation, rt.pendingLimit()), done: make(chan struct{}), } - public := &Session{state: state} + public := &Session{state: state, logicalID: logicalID} state.commands = &commandSession{state: state} mailbox.automatic = func() { state.startAutomaticRun() } - rt.sessions[id] = state + rt.sessions[logicalID] = state rt.wg.Add(1) rt.mu.Unlock() - if id == MainREPLName && rt.option != nil && rt.option.Heartbeat > 0 { + if logicalID == MainREPLName && rt.option != nil && rt.option.Heartbeat > 0 { _, _ = scheduler.Add(sessionCtx, agent.LoopEntry{ Name: "heartbeat", Interval: time.Duration(rt.option.Heartbeat) * time.Minute, Mode: agent.ModeInbox, @@ -438,6 +449,9 @@ func (rt *AgentRuntime) OpenSession(ctx context.Context, options SessionOptions) 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 { + emitContinuationMessages(state, prepareContinuationMessages(options.Messages)) + } return public, nil } @@ -449,22 +463,26 @@ func (rt *AgentRuntime) EnsureSession(options SessionOptions) (*Session, error) return nil, fmt.Errorf("agent runtime is not configured") } id := strings.TrimSpace(options.ID) - if id != "" { + logicalID := strings.TrimSpace(options.LogicalID) + if logicalID == "" { + logicalID = id + } + if logicalID != "" { rt.mu.RLock() - state := rt.sessions[id] + state := rt.sessions[logicalID] rt.mu.RUnlock() if state != nil { return ensuredSession(state, options) } } session, err := rt.OpenSession(rt.ctx, options) - if err == nil || id == "" { + if err == nil || logicalID == "" { return session, err } // Concurrent reconnects may both observe the Session as absent. The strict // OpenSession call admits one; the loser re-reads and validates that Session. rt.mu.RLock() - state := rt.sessions[id] + state := rt.sessions[logicalID] rt.mu.RUnlock() if state == nil { return nil, err @@ -482,7 +500,7 @@ func ensuredSession(state *sessionState, options SessionOptions) (*Session, erro if options.AgentName != "" && options.AgentName != state.agentName { return nil, fmt.Errorf("session %q agent name conflicts with open session", state.id) } - return &Session{state: state}, nil + return &Session{state: state, logicalID: state.logicalID}, nil } func (rt *AgentRuntime) CloseSession(ctx context.Context, sessionID string, reason SessionCloseReason) error { @@ -493,9 +511,9 @@ func (rt *AgentRuntime) CloseSession(ctx context.Context, sessionID string, reas reason = SessionCloseCompleted } rt.mu.Lock() - state := rt.sessions[sessionID] + logicalID, state := rt.findSessionLocked(sessionID) if state != nil { - delete(rt.sessions, sessionID) + delete(rt.sessions, logicalID) } rt.mu.Unlock() if state == nil { @@ -519,6 +537,19 @@ func (rt *AgentRuntime) CloseSession(ctx context.Context, sessionID string, reas return nil } +func (rt *AgentRuntime) findSessionLocked(sessionID string) (string, *sessionState) { + sessionID = strings.TrimSpace(sessionID) + if state := rt.sessions[sessionID]; state != nil { + return sessionID, state + } + for logicalID, state := range rt.sessions { + if state != nil && state.id == sessionID { + return logicalID, state + } + } + return "", nil +} + func (rt *AgentRuntime) Subscribe(fn func(*aop.Event)) func() { if rt == nil || rt.bus == nil || fn == nil { return func() {} @@ -526,17 +557,26 @@ func (rt *AgentRuntime) Subscribe(fn func(*aop.Event)) func() { return rt.bus.Subscribe(fn) } +// EmitEvent publishes an already-formed runtime event through the App-owned +// AOP bus, applying the same timestamp and sequence stamping as agent events. +func (rt *AgentRuntime) EmitEvent(event *aop.Event) { + if rt == nil || rt.sessionEvents == nil || event == nil { + return + } + rt.sessionEvents.Emit(event) +} + func (rt *AgentRuntime) session(sessionID string) (*Session, error) { if rt == nil { return nil, fmt.Errorf("agent runtime is not configured") } rt.mu.RLock() - state := rt.sessions[strings.TrimSpace(sessionID)] + logicalID, state := rt.findSessionLocked(sessionID) rt.mu.RUnlock() if state == nil { return nil, fmt.Errorf("session %q is not open", sessionID) } - return &Session{state: state}, nil + return &Session{state: state, logicalID: logicalID}, nil } func (rt *AgentRuntime) RunSession(ctx context.Context, sessionID string, input RunInput) (*Run, error) { @@ -578,8 +618,13 @@ func (rt *AgentRuntime) CancelSessionRun(sessionID, turnID string) error { turnID = strings.TrimSpace(turnID) rt.mu.RLock() run := rt.runs[turnID] + _, state := rt.findSessionLocked(sessionID) rt.mu.RUnlock() - if run == nil || run.sessionID != sessionID { + actualID := sessionID + if state != nil { + actualID = state.id + } + if run == nil || run.sessionID != actualID { return fmt.Errorf("turn %q is not active in session %q", turnID, sessionID) } run.cancel() @@ -595,28 +640,33 @@ func (rt *AgentRuntime) WaitOperations() { } func (s *Session) Run(ctx context.Context, input RunInput) (*Run, error) { - if s == nil || s.state == nil { + state := s.currentState() + if state == nil { return nil, fmt.Errorf("session is not configured") } - return s.state.startRun(ctx, input) + return state.startRun(ctx, input) } func (s *Session) Command(ctx context.Context, line string) (*types.CommandResult, error) { - if s == nil || s.state == nil { + state := s.currentState() + if state == nil { return nil, fmt.Errorf("session is not configured") } + if name := commandName(line); name == "/clear" || name == "/compact" { + return s.rotateCommand(ctx, line) + } done := make(chan commandOutcome, 1) op := &sessionOperation{ execute: func(runCtx context.Context) { - outcome := s.state.commands.execute(runCtx, line) + outcome := state.commands.execute(runCtx, line) if outcome.err == nil && len(outcome.result.GetContent()) > 0 { - s.state.emitCommandResult(outcome.result) + state.emitCommandResult(outcome.result) } done <- outcome }, reject: func(err error) { done <- commandOutcome{err: err} }, } - if err := s.state.admit(ctx, op); err != nil { + if err := state.admit(ctx, op); err != nil { return nil, err } outcome := <-done @@ -624,18 +674,246 @@ func (s *Session) Command(ctx context.Context, line string) (*types.CommandResul } func (s *Session) ID() string { - if s == nil || s.state == nil { + state := s.currentState() + if state == nil { + state = s.baseState() + } + if state == nil { return "" } - return s.state.id + return state.id } func (s *Session) MessagesSnapshot() []*aop.Message { - if s == nil || s.state == nil { + state := s.currentState() + if state == nil { return nil } - return s.state.agent.MessagesSnapshot() + return state.agent.MessagesSnapshot() +} + +func (s *Session) Agent() *agent.Agent { + state := s.currentState() + if state == nil { + state = s.baseState() + } + if state == nil { + return nil + } + return state.agent +} + +func (s *Session) currentState() *sessionState { + base := s.baseState() + if base == nil || base.runtime == nil { + return nil + } + s.mu.RLock() + logicalID := s.logicalID + s.mu.RUnlock() + if logicalID == "" { + logicalID = base.logicalID + } + base.runtime.mu.RLock() + state := base.runtime.sessions[logicalID] + base.runtime.mu.RUnlock() + return state +} + +func (s *Session) baseState() *sessionState { + if s == nil { + return nil + } + s.mu.RLock() + state := s.state + s.mu.RUnlock() + return state +} + +func commandName(line string) string { + fields := strings.Fields(strings.TrimSpace(line)) + if len(fields) == 0 { + return "" + } + return fields[0] +} + +func (s *Session) rotateCommand(ctx context.Context, line string) (*types.CommandResult, error) { + state := s.currentState() + if state == nil { + return nil, fmt.Errorf("session is not configured") + } + if state.runtime.sessionRunActive(state.id) { + return nil, fmt.Errorf("task is running — use /stop first") + } + name := commandName(line) + switch name { + case "/clear": + newState, err := s.rotate(ctx, SessionCloseCleared, state.id, nil, "") + if err != nil { + return nil, err + } + outcome := commandText(line, CommandPresentationPlain, "Context cleared.") + newState.emitCommandResult(outcome.result) + return outcome.result, nil + case "/compact": + messages := state.agent.MessagesSnapshot() + if len(messages) < 4 { + return commandText(line, CommandPresentationPlain, "Nothing to compact (too few messages).").result, nil + } + values, err := commands.SplitCommandLine(line) + if err != nil { + return nil, err + } + instructions := "" + if len(values) > 1 { + instructions = strings.TrimSpace(strings.Join(values[1:], " ")) + } + result, err := state.agent.Compact(ctx, agent.CompactConfig{CustomInstructions: instructions}) + if err != nil { + return nil, err + } + newState, err := s.rotate(ctx, SessionCloseCompacted, state.id, state.agent.MessagesSnapshot(), "") + if err != nil { + return nil, err + } + commandResult := commandText(line, CommandPresentationPlain, fmt.Sprintf( + "Compacted: ~%d -> ~%d tokens (%d messages kept)", result.TokensBefore, result.TokensAfter, result.KeptMessages)) + newState.emitCommandResult(commandResult.result) + return commandResult.result, nil + default: + return nil, fmt.Errorf("unsupported rotating command %q", name) + } +} + +func (s *Session) Resume(ctx context.Context, path string) (int, error) { + state := s.currentState() + if state == nil { + return 0, fmt.Errorf("session is not configured") + } + if state.runtime.sessionRunActive(state.id) { + return 0, fmt.Errorf("task is running — use /stop first") + } + data, err := loadResumeState(path) + if err != nil { + return 0, err + } + if err := output.ValidateJSONLTarget(path); err != nil { + return 0, err + } + if _, err := s.rotate(ctx, SessionCloseResumed, data.SessionID, data.Messages, path); err != nil { + return 0, err + } + return len(data.Messages), nil +} + +func (s *Session) rotate(ctx context.Context, reason SessionCloseReason, parentSessionID string, messages []*aop.Message, recordPath string) (*sessionState, error) { + oldState := s.currentState() + if oldState == nil { + return nil, fmt.Errorf("session is not configured") + } + rt := oldState.runtime + logicalID := oldState.logicalID + agentName := oldState.agentName + prepared := prepareContinuationMessages(messages) + if err := rt.CloseSession(ctx, logicalID, reason); err != nil { + return nil, err + } + if recordPath != "" && !samePath(rt.recordPath, recordPath) { + if rt.app == nil { + return nil, fmt.Errorf("runtime application is unavailable") + } + if err := rt.app.SwitchRecording(recordPath); err != nil { + return nil, err + } + rt.recordPath = recordPath + if rt.option != nil { + rt.option.OutputFile = recordPath + rt.option.Resume = recordPath + } + } + newID := rt.nextContinuationID(logicalID) + continuation, err := rt.OpenSession(ctx, SessionOptions{ + ID: newID, LogicalID: logicalID, ParentSessionID: parentSessionID, + AgentName: agentName, Messages: prepared, + }) + if err != nil { + return nil, err + } + newState := continuation.currentState() + if newState == nil { + return nil, fmt.Errorf("continuation session was not created") + } + s.mu.Lock() + s.state = newState + s.mu.Unlock() + return newState, nil +} + +func (rt *AgentRuntime) sessionRunActive(sessionID string) bool { + rt.mu.RLock() + defer rt.mu.RUnlock() + for _, run := range rt.runs { + if run != nil && run.sessionID == sessionID { + return true + } + } + return false +} + +func prepareContinuationMessages(messages []*aop.Message) []*aop.Message { + prepared := make([]*aop.Message, 0, len(messages)) + var counter int64 + for _, message := range messages { + if message == nil { + continue + } + cloned := proto.CloneOf(message) + counter = max(counter, continuationMessageSequence(cloned.Id)) + prepared = append(prepared, cloned) + } + for _, message := range prepared { + if strings.TrimSpace(message.Id) == "" { + counter++ + message.Id = fmt.Sprintf("m-%d", counter) + } + } + return prepared +} + +func continuationMessageSequence(id string) int64 { + if !strings.HasPrefix(id, "m-") { + return 0 + } + value, _ := strconv.ParseInt(strings.TrimPrefix(id, "m-"), 10, 64) + return value +} + +func emitContinuationMessages(state *sessionState, messages []*aop.Message) { + if state == nil || state.runtime == nil { + return + } + for _, message := range messages { + if message == nil { + continue + } + if message.Role == "tool" { + for _, content := range message.Content { + if result := content.GetToolResult(); result != nil { + state.runtime.sessionEvents.Emit(&aop.Event{ + SessionId: state.id, Emitter: state.agentName, + Payload: &aop.Event_ToolResult{ToolResult: proto.CloneOf(result)}, + }) + } + } + continue + } + state.runtime.sessionEvents.Emit(&aop.Event{ + SessionId: state.id, Emitter: state.agentName, + Payload: &aop.Event_Message{Message: proto.CloneOf(message)}, + }) + } } func (s *sessionState) startRun(ctx context.Context, input RunInput) (*Run, error) { @@ -874,6 +1152,10 @@ func (rt *AgentRuntime) nextRuntimeID(prefix string) string { return id } +func (rt *AgentRuntime) nextContinuationID(logicalID string) string { + return logicalID + "-" + rt.nextRuntimeID(fmt.Sprintf("session-%d", time.Now().UnixNano())) +} + func (rt *AgentRuntime) releaseRun(run *Run) { if run == nil { return @@ -930,5 +1212,11 @@ func (rt *AgentRuntime) consoleAppInfoForSession(session *Session) tui.AppInfo { _, err := session.Command(ctx, line) return err } + info.Resume = session.Resume + info.ListSessions = func() ([]tui.SavedSession, error) { + return listSavedSessions(cfg.DataSubDir("sessions")) + } + info.ActiveAgent = session.Agent + info.ActiveSessionID = session.ID return info } diff --git a/pkg/runner/runtime_session_test.go b/pkg/runner/runtime_session_test.go index 12e1235e..0c89eb35 100644 --- a/pkg/runner/runtime_session_test.go +++ b/pkg/runner/runtime_session_test.go @@ -4,6 +4,12 @@ import ( "context" "errors" "fmt" + "path/filepath" + "strings" + "sync" + "testing" + "time" + "github.com/chainreactors/aiscan/agent" "github.com/chainreactors/aiscan/agent/inbox" "github.com/chainreactors/aiscan/agent/provider" @@ -15,10 +21,6 @@ import ( "github.com/chainreactors/aiscan/pkg/commands" types "github.com/chainreactors/aiscan/pkg/types" "google.golang.org/protobuf/proto" - "strings" - "sync" - "testing" - "time" ) type runtimeSemanticProvider struct { @@ -382,3 +384,61 @@ func TestRuntimeSessionRejectsRequestsPastPendingLimit(t *testing.T) { t.Fatal(fmt.Errorf("empty overflow error")) } } + +func TestResolveJSONLRecordPathSemantics(t *testing.T) { + option := &cfg.Option{} + if got := resolveJSONLRecordPath(option, REPLDisabled); got != "" { + t.Fatalf("one-shot path = %q, want empty", got) + } + + option.SaveSession = true + if got := resolveJSONLRecordPath(option, REPLDisabled); !strings.HasSuffix(got, ".jsonl") || !strings.Contains(got, "sessions") { + t.Fatalf("save-session path = %q", got) + } + + option = &cfg.Option{} + if got := resolveJSONLRecordPath(option, REPLEphemeral); !strings.HasSuffix(got, ".jsonl") || !strings.Contains(got, "sessions") { + t.Fatalf("REPL path = %q", got) + } + + option.Resume = "old.jsonl" + option.SaveSession = true + if got := resolveJSONLRecordPath(option, REPLDisabled); got != "old.jsonl" { + t.Fatalf("resume path = %q", got) + } +} + +func TestRotationCommandsRejectActiveRunWithoutSwitchingSession(t *testing.T) { + target := filepath.Join(t.TempDir(), "target.jsonl") + writePersistenceSession(t, target) + provider := &runtimeSemanticProvider{started: make(chan struct{}), release: make(chan struct{})} + runtime := newBareRuntime(t, nil, provider) + session, err := runtime.OpenSession(context.Background(), SessionOptions{ID: MainREPLName}) + if err != nil { + t.Fatal(err) + } + run, err := session.Run(context.Background(), RunInput{Content: []*aop.Content{aop.Text("running")}}) + if err != nil { + t.Fatal(err) + } + <-provider.started + originalID := session.ID() + for _, command := range []string{"/clear", "/compact"} { + if _, err := session.rotateCommand(context.Background(), command); err == nil || !strings.Contains(err.Error(), "task is running") { + t.Fatalf("%s error = %v", command, err) + } + if session.ID() != originalID { + t.Fatalf("session switched during %s: %q -> %q", command, originalID, session.ID()) + } + } + if _, err := session.Resume(context.Background(), target); err == nil || !strings.Contains(err.Error(), "task is running") { + t.Fatalf("Resume error = %v", err) + } + if session.ID() != originalID { + t.Fatalf("session switched while active: %q -> %q", originalID, session.ID()) + } + close(provider.release) + if _, err := run.Wait(); err != nil { + t.Fatal(err) + } +} diff --git a/pkg/runner/session_jsonl.go b/pkg/runner/session_jsonl.go new file mode 100644 index 00000000..7eb6e681 --- /dev/null +++ b/pkg/runner/session_jsonl.go @@ -0,0 +1,139 @@ +package runner + +import ( + "fmt" + "os" + "path/filepath" + "sort" + "strconv" + "strings" + "time" + + aop "github.com/chainreactors/aiscan/aop" + "github.com/chainreactors/aiscan/core/output" + "github.com/chainreactors/aiscan/pkg/tui" + types "github.com/chainreactors/aiscan/pkg/types" + "google.golang.org/protobuf/proto" +) + +type resumeState struct { + SessionID string + Model string + Messages []*aop.Message + MessageCounter int64 +} + +type resumeStream struct { + id string + parentID string + parentToolCall string + model string + messages []*aop.Message + messageCounter int64 + order int + started bool +} + +func loadResumeState(path string) (*resumeState, error) { + streams := make(map[string]*resumeStream) + order := 0 + err := output.ScanJSONL(path, func(event *aop.Event) error { + stream := streams[event.SessionId] + if stream == nil { + order++ + stream = &resumeStream{id: event.SessionId, order: order} + streams[event.SessionId] = stream + } + switch payload := event.Payload.(type) { + case *aop.Event_SessionStarted: + stream.started = true + stream.parentID = payload.SessionStarted.ParentSessionId + stream.parentToolCall = payload.SessionStarted.ParentToolCallId + if payload.SessionStarted.Model != "" { + stream.model = payload.SessionStarted.Model + } + case *aop.Event_Message: + if payload.Message == nil || (payload.Message.Role != "user" && payload.Message.Role != "assistant") { + return nil + } + if _, command, _ := types.GetCommandDetail(event); command { + return nil + } + stream.messages = append(stream.messages, proto.CloneOf(payload.Message)) + stream.messageCounter = max(stream.messageCounter, messageIDSequence(payload.Message.Id)) + case *aop.Event_ToolResult: + if payload.ToolResult != nil { + stream.messages = append(stream.messages, &aop.Message{ + Role: "tool", Content: []*aop.Content{{Value: &aop.Content_ToolResult{ToolResult: proto.CloneOf(payload.ToolResult)}}}, + }) + } + } + return nil + }) + if err != nil { + return nil, err + } + var selected *resumeStream + for _, stream := range streams { + if !stream.started || stream.parentToolCall != "" { + continue + } + if len(stream.messages) == 0 && stream.parentID == "" && stream.model == "" { + continue + } + if selected == nil || stream.order > selected.order { + selected = stream + } + } + if selected == nil { + return nil, fmt.Errorf("no resumable AOP session found in %s", path) + } + return &resumeState{ + SessionID: selected.id, Model: selected.model, Messages: selected.messages, + MessageCounter: selected.messageCounter, + }, nil +} + +func messageIDSequence(id string) int64 { + if !strings.HasPrefix(id, "m-") { + return 0 + } + value, _ := strconv.ParseInt(strings.TrimPrefix(id, "m-"), 10, 64) + return value +} + +func listSavedSessions(dir string) ([]tui.SavedSession, error) { + entries, err := os.ReadDir(dir) + if err != nil { + if os.IsNotExist(err) { + return nil, nil + } + return nil, fmt.Errorf("read session directory: %w", err) + } + var sessions []tui.SavedSession + for _, entry := range entries { + if entry.IsDir() || !strings.EqualFold(filepath.Ext(entry.Name()), ".jsonl") { + continue + } + path := filepath.Join(dir, entry.Name()) + state, err := loadResumeState(path) + if err != nil { + continue + } + updatedAt := time.Time{} + if info, infoErr := entry.Info(); infoErr == nil { + updatedAt = info.ModTime() + } + sessions = append(sessions, tui.SavedSession{ + Path: path, SessionID: state.SessionID, Model: state.Model, + Messages: len(state.Messages), UpdatedAt: updatedAt, + }) + } + sort.Slice(sessions, func(i, j int) bool { + if sessions[i].UpdatedAt.Equal(sessions[j].UpdatedAt) { + return sessions[i].Path > sessions[j].Path + } + return sessions[i].UpdatedAt.After(sessions[j].UpdatedAt) + }) + return sessions, nil +} diff --git a/pkg/runner/session_jsonl_test.go b/pkg/runner/session_jsonl_test.go new file mode 100644 index 00000000..36de67c6 --- /dev/null +++ b/pkg/runner/session_jsonl_test.go @@ -0,0 +1,85 @@ +package runner + +import ( + "os" + "path/filepath" + "testing" + "time" + + aop "github.com/chainreactors/aiscan/aop" + toolpb "github.com/chainreactors/aiscan/aop/tool" + "github.com/chainreactors/aiscan/core/eventbus" + "github.com/chainreactors/aiscan/core/output" + "google.golang.org/protobuf/types/known/anypb" + "google.golang.org/protobuf/types/known/timestamppb" +) + +func TestLoadResumeStateRebuildsCanonicalTranscript(t *testing.T) { + path := filepath.Join(t.TempDir(), "session.jsonl") + artifact, err := anypb.New(&toolpb.Artifact{Tool: "gogo", Kind: toolpb.ArtifactKindService, Data: []byte(`{"ip":"127.0.0.1"}`), CallId: "call-1"}) + if err != nil { + t.Fatal(err) + } + writeSessionEvents(t, path, []*aop.Event{ + sessionTestEvent("root", &aop.Event{Payload: &aop.Event_SessionStarted{SessionStarted: &aop.SessionStarted{Model: "test-model"}}}), + sessionTestEvent("root", &aop.Event{Payload: &aop.Event_Message{Message: &aop.Message{Id: "m-7", Role: "user", Content: []*aop.Content{aop.Text("hello")}}}}), + sessionTestEvent("root", &aop.Event{Payload: &aop.Event_Message{Message: &aop.Message{Id: "m-8", Role: "assistant", Content: []*aop.Content{aop.Text("working")}}}}), + sessionTestEvent("root", &aop.Event{Payload: &aop.Event_ToolResult{ToolResult: &aop.ToolResult{CallId: "call-1", Name: "gogo", Output: []*aop.Content{aop.Text("done")}}}}), + sessionTestEvent("root", &aop.Event{Payload: &aop.Event_Extension{Extension: artifact}}), + sessionTestEvent("child", &aop.Event{Payload: &aop.Event_SessionStarted{SessionStarted: &aop.SessionStarted{ParentSessionId: "root", ParentToolCallId: "call-child"}}}), + sessionTestEvent("child", &aop.Event{Payload: &aop.Event_Message{Message: &aop.Message{Id: "m-99", Role: "assistant", Content: []*aop.Content{aop.Text("child")}}}}), + }) + + data, err := loadResumeState(path) + if err != nil { + t.Fatal(err) + } + if data.SessionID != "root" || data.Model != "test-model" || data.MessageCounter != 8 || len(data.Messages) != 3 { + t.Fatalf("resume state = %#v", data) + } + if result := data.Messages[2].Content[0].GetToolResult(); result == nil || result.CallId != "call-1" { + t.Fatalf("tool result message = %#v", data.Messages[2]) + } +} + +func TestListSavedSessionsOnlyReadsJSONL(t *testing.T) { + dir := t.TempDir() + writeSessionEvents(t, filepath.Join(dir, "session.jsonl"), []*aop.Event{ + sessionTestEvent("root", &aop.Event{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 := os.WriteFile(filepath.Join(dir, "legacy.json"), []byte(`{"messages":[]}`), 0o644); err != nil { + t.Fatal(err) + } + sessions, err := listSavedSessions(dir) + if err != nil { + t.Fatal(err) + } + if len(sessions) != 1 || filepath.Base(sessions[0].Path) != "session.jsonl" { + t.Fatalf("sessions = %#v", sessions) + } +} + +func sessionTestEvent(sessionID string, event *aop.Event) *aop.Event { + event.Id = "event" + event.SessionId = sessionID + event.TurnId = "turn-1" + event.Emitter = "aiscan" + event.EmittedAt = timestamppb.New(time.Date(2026, 8, 3, 0, 0, 0, 0, time.UTC)) + return event +} + +func writeSessionEvents(t *testing.T, path string, events []*aop.Event) { + t.Helper() + bus := eventbus.New[*aop.Event]() + recorder, err := output.NewJSONLRecorder(bus, path) + if err != nil { + t.Fatal(err) + } + for _, event := range events { + bus.Emit(event) + } + if err := recorder.Close(); err != nil { + t.Fatal(err) + } +} diff --git a/pkg/runner/tool_call.go b/pkg/runner/tool_call.go index bbce4681..0d721552 100644 --- a/pkg/runner/tool_call.go +++ b/pkg/runner/tool_call.go @@ -10,7 +10,6 @@ import ( aop "github.com/chainreactors/aiscan/aop" toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/tool" "github.com/chainreactors/aiscan/pkg/commands" "google.golang.org/protobuf/types/known/timestamppb" @@ -37,7 +36,7 @@ type foregroundTool interface { // ExecuteToolRequest runs one canonical AOP tool call against the executor // and wraps the outcome as a ToolResult event correlated to operationID. -func ExecuteToolRequest(ctx context.Context, operationID string, request *toolpb.Call, executor ToolExecutor, dataBus *eventbus.Bus[output.ToolDataEvent]) (*aop.Event, error) { +func ExecuteToolRequest(ctx context.Context, operationID string, request *toolpb.Call, executor ToolExecutor, progressBus *eventbus.Bus[*toolpb.Progress]) (*aop.Event, error) { if request == nil || request.Call == nil || operationID == "" { return nil, fmt.Errorf("tool call correlation is invalid") } @@ -51,12 +50,12 @@ func ExecuteToolRequest(ctx context.Context, operationID string, request *toolpb if strings.TrimSpace(call.Name) == "" { return nil, fmt.Errorf("tool name is required") } - if call.WorkingDirectory != "" { - ctx = tool.ContextWithInvocation(ctx, tool.Invocation{WorkDir: call.WorkingDirectory}) - } - ctx = output.ContextWithCallID(ctx, operationID) + ctx = tool.ContextWithInvocation(ctx, tool.Invocation{ + WorkDir: call.WorkingDirectory, CallID: operationID, + SessionID: request.SessionId, TurnID: request.TurnId, Emitter: "aiscan.agent", + }) started := time.Now() - result, execErr := executeCall(ctx, executor, call, dataBus, operationID) + result, execErr := executeCall(ctx, executor, call, progressBus, operationID) if result == nil { result = &aop.ToolResult{} } @@ -73,10 +72,9 @@ func ExecuteToolRequest(ctx context.Context, operationID string, request *toolpb }, nil } -// executeCall runs the tool call. Tools with foreground capability stream -// stdout lines as tool.data progress events on dataBus while running; all -// other tools take the plain ExecuteTool path. -func executeCall(ctx context.Context, executor ToolExecutor, call *aop.ToolCall, dataBus *eventbus.Bus[output.ToolDataEvent], callID string) (*tool.Result, error) { +// executeCall runs the tool call. Tools with foreground capability publish +// ephemeral progress while running; all other tools take the plain ExecuteTool path. +func executeCall(ctx context.Context, executor ToolExecutor, call *aop.ToolCall, progressBus *eventbus.Bus[*toolpb.Progress], callID string) (*tool.Result, error) { arguments := call.GetArguments().GetData() if len(arguments) == 0 { arguments = []byte("{}") @@ -88,7 +86,7 @@ func executeCall(ctx context.Context, executor ToolExecutor, call *aop.ToolCall, if err != nil { return nil, err } - progress := newProgressStreamer(dataBus, call.Name, callID) + progress := newProgressStreamer(progressBus, call.Name, callID) result, err := fg.RunForegroundTool(ctx, args.Command, commands.BashExecOptions{ Timeout: time.Duration(args.Timeout) * time.Second, OnOutput: progress.Write, @@ -102,9 +100,9 @@ func executeCall(ctx context.Context, executor ToolExecutor, call *aop.ToolCall, } // progressStreamer splits raw command output into lines and publishes each -// non-blank line as a tool.data progress event. +// non-blank line as ephemeral tool progress. type progressStreamer struct { - bus *eventbus.Bus[output.ToolDataEvent] + bus *eventbus.Bus[*toolpb.Progress] tool string callID string buf []byte @@ -113,7 +111,7 @@ type progressStreamer struct { // maxProgressBuf is the maximum buffer size before a progressStreamer flushes. const maxProgressBuf = 64 << 10 -func newProgressStreamer(bus *eventbus.Bus[output.ToolDataEvent], tool, callID string) *progressStreamer { +func newProgressStreamer(bus *eventbus.Bus[*toolpb.Progress], tool, callID string) *progressStreamer { return &progressStreamer{bus: bus, tool: tool, callID: callID} } @@ -149,11 +147,7 @@ func (s *progressStreamer) emit(line string) { if strings.TrimSpace(line) == "" { return } - s.bus.Emit(output.ToolDataEvent{ - Tool: s.tool, - Kind: output.ToolDataProgress, - Data: line, - CallID: s.callID, - Timestamp: time.Now(), + s.bus.Emit(&toolpb.Progress{ + Tool: s.tool, Text: line, CallId: s.callID, Timestamp: timestamppb.New(time.Now()), }) } diff --git a/pkg/runner/tool_call_test.go b/pkg/runner/tool_call_test.go index 01b0e5e8..baf95e00 100644 --- a/pkg/runner/tool_call_test.go +++ b/pkg/runner/tool_call_test.go @@ -10,7 +10,6 @@ import ( aop "github.com/chainreactors/aiscan/aop" toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/tool" "github.com/chainreactors/aiscan/pkg/commands" ) @@ -115,21 +114,19 @@ func TestExecuteToolRequestForeground(t *testing.T) { registry := commands.NewRegistry() bash := &recordingBash{} registry.RegisterTool(bash) - dataBus := eventbus.New[output.ToolDataEvent]() - var progress []output.ToolDataEvent - dataBus.Subscribe(func(event output.ToolDataEvent) { - if event.Kind == output.ToolDataProgress { - progress = append(progress, event) - } + progressBus := eventbus.New[*toolpb.Progress]() + var progress []*toolpb.Progress + progressBus.Subscribe(func(event *toolpb.Progress) { + progress = append(progress, event) }) - event, err := ExecuteToolRequest(context.Background(), "task-1", toolRequest(t, "task-1", "bash", map[string]any{"command": "echo test", "timeout": 7}), registry, dataBus) + event, err := ExecuteToolRequest(context.Background(), "task-1", toolRequest(t, "task-1", "bash", map[string]any{"command": "echo test", "timeout": 7}), registry, progressBus) if err != nil { t.Fatal(err) } if bash.command != "echo test" || bash.options.Timeout != 7*time.Second { t.Fatalf("bash options = %+v", bash.options) } - if len(progress) != 1 || progress[0].Data != "streamed" || progress[0].CallID != "task-1" { + if len(progress) != 1 || progress[0].Text != "streamed" || progress[0].CallId != "task-1" { t.Fatalf("progress = %+v", progress) } result := event.GetToolResult() diff --git a/pkg/tui/commands.go b/pkg/tui/commands.go index 5ac78a5a..b7834e9c 100644 --- a/pkg/tui/commands.go +++ b/pkg/tui/commands.go @@ -5,6 +5,7 @@ import ( "fmt" "net/url" "strings" + "time" "github.com/chainreactors/aiscan/agent" cfg "github.com/chainreactors/aiscan/core/config" @@ -25,8 +26,22 @@ type AppInfo struct { OnLoggerChange func(telemetry.Logger) Run func(context.Context, string, bool) (*agent.Result, error) Command func(context.Context, string) error + Resume func(context.Context, string) (int, error) + ListSessions func() ([]SavedSession, error) + ActiveAgent func() *agent.Agent + ActiveSessionID func() string } +type SavedSession struct { + Path string + SessionID string + Model string + Messages int + UpdatedAt time.Time +} + +func (s SavedSession) SortTime() time.Time { return s.UpdatedAt } + // Session holds the dependencies commands need to operate on. type Session struct { Ctx context.Context diff --git a/pkg/tui/console.go b/pkg/tui/console.go index daf83458..ef1b7570 100644 --- a/pkg/tui/console.go +++ b/pkg/tui/console.go @@ -461,7 +461,9 @@ func (r *AgentConsole) handleRuntimeInputLine(line string) (bool, error) { return r.appInfo.Run(ctx, prompt, false) }) } - return false, r.appInfo.Command(r.ctx, text) + err := r.appInfo.Command(r.ctx, text) + r.refreshRuntimeSession() + return false, err } func runtimeTUICommand(line string) bool { @@ -945,6 +947,20 @@ func (r *AgentConsole) ensureController() *interactiveRunController { return r.controller } +func (r *AgentConsole) refreshRuntimeSession() { + if r == nil || r.appInfo.ActiveAgent == nil { + return + } + active := r.appInfo.ActiveAgent() + if active == nil || active == r.agent { + return + } + r.agent = active + if r.controller != nil { + r.controller.SetSession(active) + } +} + func (r *AgentConsole) syncEvalToController() { if r.controller == nil { return @@ -1162,12 +1178,15 @@ func (r *AgentConsole) resumeSession(path string) error { if err != nil { return err } - data, err := agent.LoadCheckpoint(path) + if r.appInfo.Resume == nil { + return fmt.Errorf("session resume is unavailable") + } + messages, err := r.appInfo.Resume(r.ctx, path) if err != nil { return err } - r.agent.LoadMessages(data.Messages) - fmt.Fprintf(r.stdout, "Resumed %d messages from %s\n", len(data.Messages), path) + r.refreshRuntimeSession() + fmt.Fprintf(r.stdout, "Resumed %d messages from %s\n", messages, path) return nil } @@ -1192,12 +1211,11 @@ func (r *AgentConsole) renderSessions() (string, error) { return r.renderPanel("sessions", renderHelpRows(rows, colorEnabled), colorEnabled), nil } -func (r *AgentConsole) listSavedSessions() ([]agent.CheckpointInfo, error) { - dir := r.sessionDir - if dir == "" { - dir = cfg.DataSubDir("sessions") +func (r *AgentConsole) listSavedSessions() ([]SavedSession, error) { + if r.appInfo.ListSessions == nil { + return nil, fmt.Errorf("session listing is unavailable") } - return agent.ListCheckpoints(dir) + return r.appInfo.ListSessions() } func (r *AgentConsole) resumeSessionInteractive() error { @@ -1244,16 +1262,26 @@ func (r *AgentConsole) resolveSessionSelection(selector string) (string, error) if selector == "" { return "", fmt.Errorf("usage: /resume [list||#index]") } - sessions, err := r.listSavedSessions() - if err != nil { - return "", err + if strings.ContainsAny(selector, `/\`) || strings.EqualFold(filepath.Ext(selector), ".jsonl") { + return selector, nil } if idx, err := strconv.Atoi(selector); err == nil { + sessions, listErr := r.listSavedSessions() + if listErr != nil { + return "", listErr + } if idx < 1 || idx > len(sessions) { return "", fmt.Errorf("session index out of range: %d", idx) } return sessions[idx-1].Path, nil } + sessions, err := r.listSavedSessions() + if err != nil { + if r.appInfo.ListSessions == nil { + return selector, nil + } + return "", err + } for _, session := range sessions { if selector == session.Path || selector == filepath.Base(session.Path) { return session.Path, nil @@ -1262,12 +1290,12 @@ func (r *AgentConsole) resolveSessionSelection(selector string) (string, error) return selector, nil } -func sessionDetail(session agent.CheckpointInfo) string { +func sessionDetail(session SavedSession) string { parts := make([]string, 0, 4) if ts := session.SortTime(); !ts.IsZero() { parts = append(parts, ts.Local().Format("2006-01-02 15:04:05")) } - model := strings.Trim(strings.TrimSpace(session.Provider)+"/"+strings.TrimSpace(session.Model), "/") + model := strings.TrimSpace(session.Model) if model != "" { parts = append(parts, model) } diff --git a/pkg/tui/console_test.go b/pkg/tui/console_test.go index f3b1c508..c9ac2794 100644 --- a/pkg/tui/console_test.go +++ b/pkg/tui/console_test.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "encoding/json" + "fmt" "net/http" "net/http/httptest" "os" @@ -22,6 +23,7 @@ import ( "github.com/chainreactors/tui/readline/inputrc" rlterm "github.com/chainreactors/tui/readline/terminal" "google.golang.org/protobuf/encoding/protojson" + "google.golang.org/protobuf/types/known/timestamppb" ) func TestIsLocalAgentTerminal(t *testing.T) { @@ -217,6 +219,41 @@ func TestReadlineDoesNotSuppressLiveStatusWhileTaskRuns(t *testing.T) { } } +func TestAgentConsoleRefreshesAgentAfterRuntimeResumeAndClear(t *testing.T) { + var stdout, stderr bytes.Buffer + oldAgent := agent.NewAgent(agent.Config{SessionID: "old"}) + resumedAgent := agent.NewAgent(agent.Config{SessionID: "resumed"}) + clearedAgent := agent.NewAgent(agent.Config{SessionID: "cleared"}) + active := oldAgent + info := AppInfo{ + Run: func(context.Context, string, bool) (*agent.Result, error) { return &agent.Result{}, nil }, + Command: func(_ context.Context, line string) error { + if line == "/clear" { + active = clearedAgent + } + return nil + }, + Resume: func(context.Context, string) (int, error) { + active = resumedAgent + return 2, nil + }, + ActiveAgent: func() *agent.Agent { return active }, + } + repl := NewAgentConsoleWithWriters(context.Background(), &cfg.Option{}, info, oldAgent, &stdout, &stderr) + if _, err := repl.ExecuteLineAndWait("/resume session.jsonl"); err != nil { + t.Fatalf("runtime /resume: %v", err) + } + if repl.agent != resumedAgent || repl.controller.session != resumedAgent { + t.Fatal("console did not switch to the resumed runtime agent") + } + if _, err := repl.ExecuteLineAndWait("/clear"); err != nil { + t.Fatalf("runtime /clear: %v", err) + } + if repl.agent != clearedAgent || repl.controller.session != clearedAgent { + t.Fatal("console did not switch to the cleared continuation agent") + } +} + func TestAgentConsoleCtrlCWarnsAndClearsInput(t *testing.T) { var stdout, stderr bytes.Buffer repl := NewAgentConsoleWithWriters(context.Background(), &cfg.Option{}, AppInfo{}, nil, &stdout, &stderr) @@ -299,29 +336,22 @@ func TestAgentConsoleModelCommandListsAndSwitches(t *testing.T) { func TestAgentConsoleResumeLoadsSessionMessages(t *testing.T) { dir := t.TempDir() - if err := agent.SaveCheckpoint(dir, &agent.CheckpointData{ - Model: "test-model", - Provider: "capture", - Messages: []*aop.Message{ - agent.TextMessage("user", "previous user"), - agent.TextMessage("assistant", "previous assistant"), - }, - }); err != nil { - t.Fatalf("SaveSession: %v", err) - } - sessions, err := agent.ListCheckpoints(dir) - if err != nil { - t.Fatalf("ListSessions: %v", err) - } - if len(sessions) != 1 { - t.Fatalf("sessions len = %d, want 1", len(sessions)) - } - path := sessions[0].Path - + path := filepath.Join(dir, "session-resume.jsonl") + writeConsoleSession(t, path, "test-model", time.Now(), + agent.TextMessage("user", "previous user"), + agent.TextMessage("assistant", "previous assistant"), + ) var stdout, stderr bytes.Buffer prov := &captureConsoleProvider{} session := agent.NewAgent(agent.Config{Provider: prov, Model: "test-model"}) - repl := NewAgentConsoleWithWriters(context.Background(), &cfg.Option{}, AppInfo{}, session, &stdout, &stderr) + repl := NewAgentConsoleWithWriters(context.Background(), &cfg.Option{}, AppInfo{ + Resume: func(context.Context, string) (int, error) { + session.LoadMessages([]*aop.Message{ + agent.TextMessage("user", "previous user"), agent.TextMessage("assistant", "previous assistant"), + }) + return 2, nil + }, + }, session, &stdout, &stderr) if _, err := repl.ExecuteLineAndWait("/resume " + path); err != nil { t.Fatalf("/resume: %v\nstderr=%s", err, stderr.String()) @@ -348,28 +378,69 @@ func TestAgentConsoleResumeLoadsSessionMessages(t *testing.T) { t.Fatalf("request messages missing %q:\n%s", want, joined) } } + + stdout.Reset() + stderr.Reset() + if _, err := repl.ExecuteLineAndWait("/clear"); err != nil { + t.Fatalf("/clear: %v\nstderr=%s", err, stderr.String()) + } + if out := stdout.String(); !strings.Contains(out, "Context cleared.") { + t.Fatalf("clear output = %q", out) + } + if messages := session.MessagesSnapshot(); len(messages) != 0 { + t.Fatalf("messages after clear = %d, want 0", len(messages)) + } + + stdout.Reset() + stderr.Reset() + if _, err := repl.ExecuteLineAndWait("after clear"); err != nil { + t.Fatalf("prompt after clear: %v\nstderr=%s", err, stderr.String()) + } + if len(prov.requests) != 2 { + t.Fatalf("provider requests = %d, want 2", len(prov.requests)) + } + var afterClear []string + for _, msg := range prov.requests[1].Messages { + afterClear = append(afterClear, provider.MessageText(msg)) + } + afterClearText := strings.Join(afterClear, "\n") + if !strings.Contains(afterClearText, "after clear") { + t.Fatalf("request after clear missing new prompt:\n%s", afterClearText) + } + for _, stale := range []string{"previous user", "previous assistant", "new prompt"} { + if strings.Contains(afterClearText, stale) { + t.Fatalf("request after clear retained %q:\n%s", stale, afterClearText) + } + } } func TestAgentConsoleResumeListsAndSelectsSession(t *testing.T) { dir := t.TempDir() - oldPath := filepath.Join(dir, "session-old.json") - newPath := filepath.Join(dir, "session-new.json") - writeConsoleSession(t, oldPath, "old-model", "old message", time.Date(2026, 7, 12, 10, 0, 0, 0, time.UTC)) - writeConsoleSession(t, newPath, "new-model", "new message", time.Date(2026, 7, 13, 10, 0, 0, 0, time.UTC)) + oldPath := filepath.Join(dir, "session-old.jsonl") + newPath := filepath.Join(dir, "session-new.jsonl") + writeConsoleSession(t, oldPath, "old-model", time.Date(2026, 7, 12, 10, 0, 0, 0, time.UTC), agent.TextMessage("user", "old message")) + writeConsoleSession(t, newPath, "new-model", time.Date(2026, 7, 13, 10, 0, 0, 0, time.UTC), agent.TextMessage("user", "new message")) var stdout, stderr bytes.Buffer session := agent.NewAgent(agent.Config{}) - repl := NewAgentConsoleWithWriters(context.Background(), &cfg.Option{}, AppInfo{}, session, &stdout, &stderr) + saved := []SavedSession{ + {Path: newPath, Model: "new-model", Messages: 1, UpdatedAt: time.Date(2026, 7, 13, 10, 0, 0, 0, time.UTC)}, + {Path: oldPath, Model: "old-model", Messages: 1, UpdatedAt: time.Date(2026, 7, 12, 10, 0, 0, 0, time.UTC)}, + } + repl := NewAgentConsoleWithWriters(context.Background(), &cfg.Option{}, AppInfo{ + ListSessions: func() ([]SavedSession, error) { return saved, nil }, + Resume: func(context.Context, string) (int, error) { return 1, nil }, + }, session, &stdout, &stderr) repl.sessionDir = dir if _, err := repl.ExecuteLineAndWait("/resume list"); err != nil { t.Fatalf("/resume list: %v\nstderr=%s", err, stderr.String()) } listOut := stdout.String() - if !strings.Contains(listOut, "session-new.json") || !strings.Contains(listOut, "session-old.json") { + if !strings.Contains(listOut, "session-new.jsonl") || !strings.Contains(listOut, "session-old.jsonl") { t.Fatalf("resume list missing sessions:\n%s", listOut) } - if strings.Index(listOut, "session-new.json") > strings.Index(listOut, "session-old.json") { + if strings.Index(listOut, "session-new.jsonl") > strings.Index(listOut, "session-old.jsonl") { t.Fatalf("sessions not sorted newest first:\n%s", listOut) } @@ -383,22 +454,33 @@ func TestAgentConsoleResumeListsAndSelectsSession(t *testing.T) { } } -func writeConsoleSession(t *testing.T, path, model, content string, updatedAt time.Time) { +func writeConsoleSession(t *testing.T, path, model string, updatedAt time.Time, messages ...*aop.Message) { t.Helper() - msgRaw, err := protojson.Marshal(agent.TextMessage("user", content)) - if err != nil { - t.Fatalf("marshal message: %v", err) + events := []*aop.Event{{ + Id: "e-1", SessionId: "console-session", Emitter: "aiscan", EmittedAt: timestamppb.New(updatedAt), + Payload: &aop.Event_SessionStarted{SessionStarted: &aop.SessionStarted{Model: model}}, + }} + for i, message := range messages { + message.Id = fmt.Sprintf("m-%d", i+1) + events = append(events, &aop.Event{ + Id: fmt.Sprintf("e-%d", i+2), SessionId: "console-session", TurnId: "turn-1", + Emitter: "aiscan", EmittedAt: timestamppb.New(updatedAt), + Payload: &aop.Event_Message{Message: message}, + }) } - raw, err := json.Marshal(map[string]any{ - "version": 1, - "updated_at": updatedAt, - "model": model, - "messages": []json.RawMessage{msgRaw}, - }) - if err != nil { - t.Fatalf("marshal session: %v", err) + var lines strings.Builder + for _, event := range events { + raw, err := protojson.Marshal(event) + if err != nil { + t.Fatalf("marshal session event: %v", err) + } + lines.Write(raw) + lines.WriteByte('\n') } - if err := os.WriteFile(path, raw, 0o644); err != nil { + if err := os.WriteFile(path, []byte(lines.String()), 0o644); err != nil { t.Fatalf("write session: %v", err) } + if err := os.Chtimes(path, updatedAt, updatedAt); err != nil { + t.Fatalf("set session time: %v", err) + } } diff --git a/pkg/tui/controller.go b/pkg/tui/controller.go index d0d1b998..1f7272a3 100644 --- a/pkg/tui/controller.go +++ b/pkg/tui/controller.go @@ -214,6 +214,15 @@ func (c *interactiveRunController) SetOnFinish(fn func()) { c.onFinish = fn } +func (c *interactiveRunController) SetSession(session *agent.Agent) { + if c == nil || session == nil { + return + } + c.mu.Lock() + c.session = session + c.mu.Unlock() +} + func (c *interactiveRunController) notifyFinish() { c.mu.Lock() fn := c.onFinish diff --git a/pkg/tui/remote_console.go b/pkg/tui/remote_console.go index f2f2dd82..a76cacba 100644 --- a/pkg/tui/remote_console.go +++ b/pkg/tui/remote_console.go @@ -5,6 +5,7 @@ import ( "context" "fmt" "io" + "strings" "sync" "github.com/chainreactors/aiscan/agent" @@ -35,7 +36,7 @@ func RunAgentConsoleWithTerminal(ctx context.Context, option *cfg.Option, appInf return fmt.Errorf("terminal is nil") } agentOutput := NewAgentOutputWithWriters(option, terminal.Out, terminal.Err, terminal.Control == nil || terminal.Control.IsTerminal()) - unsubscribe := subscribeAgentOutput(agentOutput, session, subscribers...) + unsubscribe := subscribeAgentOutput(agentOutput, appInfo, session, subscribers...) defer unsubscribe() repl := NewAgentConsoleWithTerminal(ctx, option, appInfo, session, agentOutput, terminal) return repl.Start() @@ -43,18 +44,31 @@ func RunAgentConsoleWithTerminal(ctx context.Context, option *cfg.Option, appInf // subscribeAgentOutput filters the shared runtime bus by session ID so a // remote or local REPL cannot render sibling/subagent events accidentally. -func subscribeAgentOutput(output *AgentOutput, session *agent.Agent, subscribers ...AOPEventSubscriber) func() { +func subscribeAgentOutput(output *AgentOutput, appInfo AppInfo, session *agent.Agent, subscribers ...AOPEventSubscriber) func() { if output == nil || session == nil || len(subscribers) == 0 || subscribers[0] == nil { return func() {} } - sessionID := session.SessionID() return subscribers[0](func(event *aop.Event) { - if sessionID == "" || event.SessionId == sessionID { + sessionID := session.SessionID() + if appInfo.ActiveSessionID != nil { + sessionID = appInfo.ActiveSessionID() + } + if (sessionID == "" || event.SessionId == sessionID) && !isSessionBootstrapEvent(event) { output.HandleEvent(event) } }) } +func isSessionBootstrapEvent(event *aop.Event) bool { + if event == nil || event.TurnId != "" { + return false + } + if message := event.GetMessage(); message != nil { + return strings.HasPrefix(message.Id, "m-") + } + return event.GetToolResult() != nil +} + type remoteTerminalWriter struct { mu sync.Mutex w io.Writer diff --git a/pkg/tui/remote_console_test.go b/pkg/tui/remote_console_test.go index 49bdac08..bcd2f151 100644 --- a/pkg/tui/remote_console_test.go +++ b/pkg/tui/remote_console_test.go @@ -17,7 +17,7 @@ func TestSubscribeAgentOutputRestoresSessionEvents(t *testing.T) { session := agent.NewAgent(agent.Config{SessionID: "main-repl"}) var handler func(*aop.Event) unsubscribed := false - unsubscribe := subscribeAgentOutput(output, session, func(fn func(*aop.Event)) func() { + unsubscribe := subscribeAgentOutput(output, AppInfo{}, session, func(fn func(*aop.Event)) func() { handler = fn return func() { unsubscribed = true } }) @@ -41,3 +41,46 @@ func TestSubscribeAgentOutputRestoresSessionEvents(t *testing.T) { t.Fatal("event subscription was not released") } } + +func TestSubscribeAgentOutputTracksRotatedRuntimeSession(t *testing.T) { + var stdout bytes.Buffer + var stderr syncedBuffer + output := NewAgentOutputWithWriters(nil, &stdout, &stderr, true) + defer output.live.Stop() + + activeID := "session-old" + session := agent.NewAgent(agent.Config{SessionID: activeID}) + var handler func(*aop.Event) + unsubscribe := subscribeAgentOutput(output, AppInfo{ActiveSessionID: func() string { return activeID }}, session, func(fn func(*aop.Event)) func() { + handler = fn + return func() {} + }) + defer unsubscribe() + + activeID = "session-new" + old := turnStartEvent(1) + old.SessionId = "session-old" + handler(old) + if liveRunning(output.live) { + t.Fatal("output consumed an event from the rotated-out session") + } + current := turnStartEvent(1) + current.SessionId = "session-new" + handler(current) + if !liveRunning(output.live) { + t.Fatal("output did not follow the rotated runtime session") + } +} + +func TestSessionBootstrapEventsAreNotRenderedAsLiveOutput(t *testing.T) { + bootstrap := &aop.Event{SessionId: "next", Payload: &aop.Event_Message{Message: &aop.Message{ + Id: "m-4", Role: "assistant", Content: []*aop.Content{aop.Text("restored history")}, + }}} + if !isSessionBootstrapEvent(bootstrap) { + t.Fatal("restored message was not recognized as a bootstrap event") + } + bootstrap.TurnId = "turn-1" + if isSessionBootstrapEvent(bootstrap) { + t.Fatal("live turn message was mistaken for bootstrap history") + } +} diff --git a/pkg/web/service/agents_mux.go b/pkg/web/service/agents_mux.go index 5aed6b0e..951e019d 100644 --- a/pkg/web/service/agents_mux.go +++ b/pkg/web/service/agents_mux.go @@ -2,9 +2,7 @@ package service import ( "context" - "encoding/json" "fmt" - "time" aop "github.com/chainreactors/aiscan/aop" execpb "github.com/chainreactors/aiscan/aop/exec" @@ -267,14 +265,11 @@ func (p *AgentPool) handleToolArtifact(ctx context.Context, envelope *aop.Envelo if operationID == "" { operationID = envelope.Id } - artifact := output.ToolArtifact{ - Tool: value.Tool, Kind: value.Kind, Target: value.Target, - Data: append(json.RawMessage(nil), value.Data...), CallID: operationID, + if value.CallId == "" { + value = protobuf.CloneOf(value) + value.CallId = operationID } - if value.Timestamp != nil { - artifact.Timestamp = value.Timestamp.AsTime() - } - _ = p.artifacts.IngestArtifact(ctx, operationID, artifact) + _ = p.artifacts.IngestArtifact(ctx, value) } func (p *AgentPool) finishAgentTask(agent *remoteAgent, taskID string, result taskResult) { @@ -316,6 +311,15 @@ func (p *AgentPool) forwardAOPFrame(agent *remoteAgent, correlationID string, ev p.sessions.BroadcastAOPEvent(sessionID, event) } } + if extension := event.GetExtension(); extension != nil && p.artifacts != nil { + artifact := new(toolpb.Artifact) + if extension.MessageIs(artifact) && extension.UnmarshalTo(artifact) == nil { + if artifact.CallId == "" { + artifact.CallId = correlationID + } + _ = p.artifacts.IngestArtifact(context.Background(), artifact) + } + } switch event.Payload.(type) { case *aop.Event_TurnEnded: agent.state().convergeOnTurnEnd(event.TurnId, event) @@ -325,25 +329,18 @@ func (p *AgentPool) forwardAOPFrame(agent *remoteAgent, correlationID string, ev } func (p *AgentPool) handleToolProgress(operationID string, value *toolpb.Progress) { - if value == nil { + if value == nil || p.hub == nil { return } - event := output.ToolDataEvent{Kind: output.ToolDataProgress, Data: value.Text, CallID: operationID} - if value.Timestamp != nil { - event.Timestamp = value.Timestamp.AsTime() - } else { - event.Timestamp = time.Now() + if value.CallId != "" { + operationID = value.CallId } - if event.Kind != output.ToolDataProgress || p.hub == nil || event.CallID == "" { - return - } - line, ok := event.Data.(string) - if !ok { + if operationID == "" { return } - line = output.StripANSI(line) + line := output.StripANSI(value.Text) if line != "" { - p.hub.BroadcastScan(managementapi.ScanProgressEvent(event.CallID, line), false) + p.hub.BroadcastScan(managementapi.ScanProgressEvent(operationID, line), false) } } diff --git a/pkg/web/service/agents_test.go b/pkg/web/service/agents_test.go index 010cbbbd..64465752 100644 --- a/pkg/web/service/agents_test.go +++ b/pkg/web/service/agents_test.go @@ -6,7 +6,6 @@ import ( filepb "github.com/chainreactors/aiscan/aop/file" ptypb "github.com/chainreactors/aiscan/aop/pty" toolpb "github.com/chainreactors/aiscan/aop/tool" - "github.com/chainreactors/aiscan/core/output" types "github.com/chainreactors/aiscan/pkg/types" webstatic "github.com/chainreactors/aiscan/web" "github.com/go-rod/rod" @@ -122,12 +121,11 @@ func ptyMessageKind(value *ptypb.ProtocolMessage) string { } type recordingArtifactSink struct { - operationID string - artifact output.ToolArtifact + artifact *toolpb.Artifact } -func (s *recordingArtifactSink) IngestArtifact(_ context.Context, operationID string, artifact output.ToolArtifact) error { - s.operationID, s.artifact = operationID, artifact +func (s *recordingArtifactSink) IngestArtifact(_ context.Context, artifact *toolpb.Artifact) error { + s.artifact = protobuf.CloneOf(artifact) return nil } @@ -144,11 +142,11 @@ func TestAgentPoolForwardsRawToolArtifact(t *testing.T) { pool.SetArtifactIngestor(sink) raw := []byte(`{"ip":"127.0.0.1","port":"80"}`) pool.handleAgentEnvelope(&remoteAgent{nodeState: newNodeState()}, wrapMessage(t, generateID(), "call-gogo-1", &toolpb.ProtocolMessage{Message: &toolpb.ProtocolMessage_Artifact{Artifact: &toolpb.Artifact{ - Tool: "gogo", Kind: output.ToolDataService, Data: raw, MediaType: aop.JSONMediaType, + Tool: "gogo", Kind: toolpb.ArtifactKindService, Data: raw, MediaType: aop.JSONMediaType, }}})) - if sink.operationID != "call-gogo-1" { - t.Fatalf("operation id = %q, want tool call id", sink.operationID) + if sink.artifact.CallId != "call-gogo-1" { + t.Fatalf("operation id = %q, want tool call id", sink.artifact.CallId) } if sink.artifact.Tool != "gogo" || string(sink.artifact.Data) != string(raw) { t.Fatalf("forwarded artifact = %+v", sink.artifact) diff --git a/pkg/web/service/artifacts.go b/pkg/web/service/artifacts.go index 561eab96..4eb49ed9 100644 --- a/pkg/web/service/artifacts.go +++ b/pkg/web/service/artifacts.go @@ -3,13 +3,13 @@ package service import ( "context" - "github.com/chainreactors/aiscan/core/output" + toolpb "github.com/chainreactors/aiscan/aop/tool" ) // ArtifactIngestor is the server-side normalization boundary. Agent nodes send // scanner-native records; implementations convert and persist canonical SCO. type ArtifactIngestor interface { - IngestArtifact(context.Context, string, output.ToolArtifact) error + IngestArtifact(context.Context, *toolpb.Artifact) error NormalizeArtifact(context.Context, string, string, []byte) (uint64, uint64, error) SupportedArtifacts() []string Close() error diff --git a/pkg/web/service/artifacts_native.go b/pkg/web/service/artifacts_native.go index 62880c8e..fc56ae10 100644 --- a/pkg/web/service/artifacts_native.go +++ b/pkg/web/service/artifacts_native.go @@ -7,7 +7,7 @@ import ( "strings" "sync" - "github.com/chainreactors/aiscan/core/output" + toolpb "github.com/chainreactors/aiscan/aop/tool" cstx "github.com/chainreactors/libcstx/go" ) @@ -38,8 +38,11 @@ func NewArtifactIngestor(store SCOStore) (ArtifactIngestor, error) { return &cstxArtifactIngestor{store: store, runtime: runtime, artifacts: artifacts}, nil } -func (i *cstxArtifactIngestor) IngestArtifact(ctx context.Context, operationID string, artifact output.ToolArtifact) error { - _, _, err := i.NormalizeArtifact(ctx, operationID, artifact.Tool, artifact.Data) +func (i *cstxArtifactIngestor) IngestArtifact(ctx context.Context, artifact *toolpb.Artifact) error { + if artifact == nil { + return nil + } + _, _, err := i.NormalizeArtifact(ctx, artifact.CallId, artifact.Tool, artifact.Data) return err } diff --git a/pkg/web/service/artifacts_test.go b/pkg/web/service/artifacts_test.go index a53b9449..cadec557 100644 --- a/pkg/web/service/artifacts_test.go +++ b/pkg/web/service/artifacts_test.go @@ -5,7 +5,7 @@ import ( "encoding/json" "testing" - "github.com/chainreactors/aiscan/core/output" + toolpb "github.com/chainreactors/aiscan/aop/tool" ) type artifactTestStore struct { @@ -27,9 +27,9 @@ func TestCSTXArtifactIngestorNormalizesOnServer(t *testing.T) { } t.Cleanup(func() { _ = ingestor.Close() }) - err = ingestor.IngestArtifact(context.Background(), "scan-1", output.ToolArtifact{ - Tool: "gogo", - Data: json.RawMessage(`{"ip":"192.0.2.1","port":"80","protocol":"tcp","status":"200","uri":"http://192.0.2.1/","title":"Test"}`), + err = ingestor.IngestArtifact(context.Background(), &toolpb.Artifact{ + CallId: "scan-1", Tool: "gogo", + Data: []byte(`{"ip":"192.0.2.1","port":"80","protocol":"tcp","status":"200","uri":"http://192.0.2.1/","title":"Test"}`), }) if err != nil { t.Fatal(err) diff --git a/pkg/web/service/scan.go b/pkg/web/service/scan.go index 23ea1535..a30da281 100644 --- a/pkg/web/service/scan.go +++ b/pkg/web/service/scan.go @@ -14,6 +14,7 @@ import ( aop "github.com/chainreactors/aiscan/aop" "github.com/chainreactors/aiscan/core/output" + coretool "github.com/chainreactors/aiscan/core/tool" "github.com/chainreactors/aiscan/pkg/commands" types "github.com/chainreactors/aiscan/pkg/types" managementapi "github.com/chainreactors/aiscan/pkg/web/api" @@ -260,7 +261,7 @@ func (s *Service) runScanViaAgent(ctx context.Context, scan *types.Scan) { } func (s *Service) runScanLocally(ctx context.Context, scan *types.Scan) { - ctx = output.ContextWithCallID(ctx, scan.Id) + ctx = coretool.ContextWithInvocation(ctx, coretool.Invocation{CallID: scan.Id, Emitter: "scan"}) streamWriter := &scanStreamWriter{ hub: s.hub, scanID: scan.Id, diff --git a/tools/gogo/gogo.go b/tools/gogo/gogo.go index bfa07f4a..7820e997 100644 --- a/tools/gogo/gogo.go +++ b/tools/gogo/gogo.go @@ -7,8 +7,8 @@ import ( "path/filepath" "strings" - "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/core/output" + aop "github.com/chainreactors/aiscan/aop" + toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/tools/toolargs" @@ -38,8 +38,8 @@ func (c *Command) WithProxy(proxy string) *Command { return c } -func (c *Command) WithDataBus(bus *eventbus.Bus[output.ToolDataEvent]) *Command { - c.DataBus = bus +func (c *Command) WithEvents(events aop.EventEmitter) *Command { + c.Events = events return c } @@ -92,7 +92,7 @@ func (c *Command) Run(ctx context.Context, execution *commands.Execution) (_ any return c.engine.Init() }, OnResult: func(r *parsers.GOGOResult) { - c.EmitDataCtx(ctx, "gogo", output.ToolDataService, r.GetTarget(), r) + c.EmitArtifactCtx(ctx, "gogo", toolpb.ArtifactKindService, r.GetTarget(), r) }, } if err := gogocore.RunWithArgs(ctx, args, opts); err != nil { diff --git a/tools/gogo/register.go b/tools/gogo/register.go index d212ded8..10bed77f 100644 --- a/tools/gogo/register.go +++ b/tools/gogo/register.go @@ -21,7 +21,7 @@ func init() { d.Skip("gogo", deps.Name(engine.SetKey)+".Gogo") return } - impl := New(es.Gogo).WithLogger(d.GetLogger()).WithProxy(d.ScannerProxy).WithDataBus(d.DataBus) + impl := New(es.Gogo).WithLogger(d.GetLogger()).WithProxy(d.ScannerProxy).WithEvents(d.Events) reg.Register(commands.Command{ Name: impl.Name(), Usage: impl.Usage(), QuickReference: impl.QuickReference(), DescriptionPath: "aiscan://skills/aiscan/okf/easm/gogo.md", diff --git a/tools/katana/katana.go b/tools/katana/katana.go index a6cd5ce6..daa850dc 100644 --- a/tools/katana/katana.go +++ b/tools/katana/katana.go @@ -11,8 +11,8 @@ import ( "sync" "time" - "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/core/output" + aop "github.com/chainreactors/aiscan/aop" + toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/tools/toolargs" @@ -51,8 +51,8 @@ func (c *Command) WithProxy(proxy string) *Command { return c } -func (c *Command) WithDataBus(bus *eventbus.Bus[output.ToolDataEvent]) *Command { - c.DataBus = bus +func (c *Command) WithEvents(events aop.EventEmitter) *Command { + c.Events = events return c } @@ -162,7 +162,7 @@ func (c *Command) Run(ctx context.Context, execution *commands.Execution) (_ any options.OnResult = func(r katanaoutput.Result) { collector.collect(&r) if r.Request != nil && r.Request.URL != "" { - c.EmitDataCtx(ctx, "katana", output.ToolDataWeb, r.Request.URL, &r) + c.EmitArtifactCtx(ctx, "katana", toolpb.ArtifactKindWeb, r.Request.URL, &r) } } diff --git a/tools/katana/register.go b/tools/katana/register.go index 8fd6b2cd..3c666296 100644 --- a/tools/katana/register.go +++ b/tools/katana/register.go @@ -13,7 +13,7 @@ func init() { Capability: "katana", Build: func(deps *commands.Deps, reg *commands.CommandRegistry) { logger := deps.GetLogger() - impl := New().WithLogger(logger).WithProxy(deps.ScannerProxy).WithDataBus(deps.DataBus) + impl := New().WithLogger(logger).WithProxy(deps.ScannerProxy).WithEvents(deps.Events) reg.Register(commands.Command{ Name: impl.Name(), Usage: impl.Usage(), DescriptionPath: "aiscan://skills/aiscan/okf/easm/katana.md", diff --git a/tools/neutron/neutron.go b/tools/neutron/neutron.go index 53f675f2..ba6d1ddb 100644 --- a/tools/neutron/neutron.go +++ b/tools/neutron/neutron.go @@ -12,8 +12,8 @@ import ( "strings" "time" - "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/core/output" + aop "github.com/chainreactors/aiscan/aop" + toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" scanengine "github.com/chainreactors/aiscan/tools/scan/engine" @@ -101,8 +101,8 @@ func (c *Command) WithProxy(proxy string) *Command { return c } -func (c *Command) WithDataBus(bus *eventbus.Bus[output.ToolDataEvent]) *Command { - c.DataBus = bus +func (c *Command) WithEvents(events aop.EventEmitter) *Command { + c.Events = events return c } @@ -249,7 +249,7 @@ func (c *Command) Run(ctx context.Context, execution *commands.Execution) (_ any results = append(results, result.TemplateResult(target)) if record.Matched { summary.Matched++ - c.EmitDataCtx(ctx, "neutron", output.ToolDataVuln, target, &record) + c.EmitArtifactCtx(ctx, "neutron", toolpb.ArtifactKindVuln, target, &record) } if shouldPrintNeutronResult(record, flags) { line := formatNeutronResult(record, jsonOutput) diff --git a/tools/neutron/register.go b/tools/neutron/register.go index 3da6b4ca..655d2ac5 100644 --- a/tools/neutron/register.go +++ b/tools/neutron/register.go @@ -21,7 +21,7 @@ func init() { d.Skip("neutron", deps.Name(engine.SetKey)+".Neutron") return } - impl := New(es.Neutron, es.Index).WithLogger(d.GetLogger()).WithProxy(d.ScannerProxy).WithDataBus(d.DataBus) + impl := New(es.Neutron, es.Index).WithLogger(d.GetLogger()).WithProxy(d.ScannerProxy).WithEvents(d.Events) reg.Register(commands.Command{ Name: impl.Name(), Usage: impl.Usage(), QuickReference: impl.QuickReference(), DescriptionPath: "aiscan://skills/aiscan/okf/easm/neutron.md", diff --git a/tools/proton/command.go b/tools/proton/command.go index 210fbc55..f1090d21 100644 --- a/tools/proton/command.go +++ b/tools/proton/command.go @@ -15,8 +15,8 @@ import ( "sync/atomic" "time" - "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/core/output" + aop "github.com/chainreactors/aiscan/aop" + toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/tools/toolargs" @@ -48,8 +48,8 @@ func (c *Command) WithProxy(proxy string) *Command { return c } -func (c *Command) WithDataBus(bus *eventbus.Bus[output.ToolDataEvent]) *Command { - c.DataBus = bus +func (c *Command) WithEvents(events aop.EventEmitter) *Command { + c.Events = events return c } @@ -257,7 +257,7 @@ func (c *Command) Run(ctx context.Context, execution *commands.Execution) (_ any if uf.Class == "extract" { atomic.AddInt64(&extractCount, 1) } - c.EmitDataCtx(ctx, "proton", output.ToolDataVuln, uf.FilePath, &uf) + c.EmitArtifactCtx(ctx, "proton", toolpb.ArtifactKindVuln, uf.FilePath, &uf) writeFinding(execution.Stdout, uf, flags.JSON, inputs[0]) if fileOut != nil { writeFinding(fileOut, uf, flags.JSON, inputs[0]) diff --git a/tools/proton/register.go b/tools/proton/register.go index b3459cbb..ea2d3c91 100644 --- a/tools/proton/register.go +++ b/tools/proton/register.go @@ -16,7 +16,7 @@ func init() { commands.RegisterFactory(commands.Factory{ Capability: "proton", Build: func(d *commands.Deps, reg *commands.CommandRegistry) { - cmd := New().WithLogger(d.GetLogger()).WithProxy(d.ScannerProxy).WithDataBus(d.DataBus) + cmd := New().WithLogger(d.GetLogger()).WithProxy(d.ScannerProxy).WithEvents(d.Events) if rs, ok := deps.Get(d.Bag, resources.SetKey); ok && rs != nil { cmd.WithResourceProvider(rs.ProtonConfig) } else { diff --git a/tools/register_command.go b/tools/register_command.go index f2be5c97..b674c168 100644 --- a/tools/register_command.go +++ b/tools/register_command.go @@ -29,8 +29,8 @@ func init() { if d.ScannerProxy != "" { scanOpts = append(scanOpts, scan.WithProxy(d.ScannerProxy)) } - if d.DataBus != nil { - scanOpts = append(scanOpts, scan.WithDataBus(d.DataBus)) + if d.Events != nil { + scanOpts = append(scanOpts, scan.WithEvents(d.Events)) } impl := scan.New(es, scanOpts...) diff --git a/tools/register_command_full_integration_test.go b/tools/register_command_full_integration_test.go index 52e1599e..987df967 100644 --- a/tools/register_command_full_integration_test.go +++ b/tools/register_command_full_integration_test.go @@ -13,6 +13,8 @@ import ( "testing" "time" + aop "github.com/chainreactors/aiscan/aop" + toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/capability" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" @@ -107,10 +109,10 @@ func TestFullScannerPublicIntegration(t *testing.T) { t.Skip("set AISCAN_INTEGRATION=1 to run public network regression tests") } - bus := eventbus.New[output.ToolDataEvent]() + bus := eventbus.New[*aop.Event]() recorder := newFunctionalRecorder(bus) registry := commands.NewRegistry() - deps := &commands.Deps{WorkDir: t.TempDir(), DataBus: bus, Logger: telemetry.NopLogger()} + deps := &commands.Deps{WorkDir: t.TempDir(), Events: bus, Logger: telemetry.NopLogger()} commands.Provide(deps, engine.SetKey, &engine.Set{}) commands.BuildPlan(capability.Select(capability.Options{Groups: []string{"scanner"}}), deps, registry) @@ -123,7 +125,7 @@ func TestFullScannerPublicIntegration(t *testing.T) { Timeout: 45 * time.Second, Check: func(t *testing.T, result functionalResult) { requireOutputContains(t, result, "https://redhaze.top") - requireEvent(t, result, "katana", output.ToolDataWeb, func(data any) bool { + requireEvent(t, result, "katana", toolpb.ArtifactKindWeb, func(data any) bool { encoded, err := json.Marshal(data) return err == nil && strings.Contains(string(encoded), "redhaze.top") }) diff --git a/tools/register_command_full_test.go b/tools/register_command_full_test.go index 19efc62c..c8bdc24a 100644 --- a/tools/register_command_full_test.go +++ b/tools/register_command_full_test.go @@ -9,9 +9,10 @@ import ( "testing" "time" + aop "github.com/chainreactors/aiscan/aop" + toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/capability" "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" _ "github.com/chainreactors/aiscan/tools/katana" @@ -51,13 +52,13 @@ func TestRegisterAllRegistersPassiveWithUncover(t *testing.T) { func TestFullScannerFunctionalRegression(t *testing.T) { httpServer := newScannerHTTPFixture(t) - bus := eventbus.New[output.ToolDataEvent]() + bus := eventbus.New[*aop.Event]() recorder := newFunctionalRecorder(bus) registry := commands.NewRegistry() engineSet := &engine.Set{} deps := &commands.Deps{ WorkDir: t.TempDir(), - DataBus: bus, + Events: bus, Logger: telemetry.NopLogger(), } commands.Provide(deps, engine.SetKey, engineSet) @@ -80,7 +81,7 @@ func TestFullScannerFunctionalRegression(t *testing.T) { Timeout: 30 * time.Second, Check: func(t *testing.T, result functionalResult) { requireOutputContains(t, result, "/admin", "/app.js", "/api/status") - requireEvent(t, result, "katana", output.ToolDataWeb, func(data any) bool { + requireEvent(t, result, "katana", toolpb.ArtifactKindWeb, func(data any) bool { encoded, err := json.Marshal(data) return err == nil && strings.Contains(string(encoded), "/admin") }) diff --git a/tools/register_command_integration_test.go b/tools/register_command_integration_test.go index 3dbe9689..f368eb38 100644 --- a/tools/register_command_integration_test.go +++ b/tools/register_command_integration_test.go @@ -10,6 +10,8 @@ import ( "testing" "time" + aop "github.com/chainreactors/aiscan/aop" + toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/capability" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" @@ -36,12 +38,12 @@ func TestScannerPublicIntegration(t *testing.T) { } defer engineSet.Close() - bus := eventbus.New[output.ToolDataEvent]() + bus := eventbus.New[*aop.Event]() recorder := newFunctionalRecorder(bus) registry := commands.NewRegistry() deps := &commands.Deps{ WorkDir: t.TempDir(), - DataBus: bus, Logger: telemetry.NopLogger(), + Events: bus, Logger: telemetry.NopLogger(), } commands.Provide(deps, engine.SetKey, engineSet) commands.Provide(deps, resources.SetKey, engineSet.Resources) @@ -69,7 +71,7 @@ http: Timeout: 45 * time.Second, Check: func(t *testing.T, result functionalResult) { requireOutputContains(t, result, `"port":"80"`, `"port":"443"`, "nginx") - requireEvent(t, result, "gogo", output.ToolDataService, func(data any) bool { + requireEvent(t, result, "gogo", toolpb.ArtifactKindService, func(data any) bool { item, ok := data.(*parsers.GOGOResult) return ok && item != nil && item.Port == "443" && item.Protocol == "https" }) @@ -96,7 +98,7 @@ http: Timeout: 30 * time.Second, Check: func(t *testing.T, result functionalResult) { requireOutputContains(t, result, `"matched":true`, `"template":"redhaze-public-marker"`) - requireEvent(t, result, "neutron", output.ToolDataVuln, nil) + requireEvent(t, result, "neutron", toolpb.ArtifactKindVuln, nil) }, }, { @@ -108,7 +110,7 @@ http: Timeout: 90 * time.Second, Check: func(t *testing.T, result functionalResult) { requireOutputContains(t, result, "[summary] completed", "443", "nginx") - requireEvent(t, result, "gogo", output.ToolDataService, func(data any) bool { + requireEvent(t, result, "gogo", toolpb.ArtifactKindService, func(data any) bool { item, ok := data.(*parsers.GOGOResult) return ok && item != nil && item.Port == "443" && item.Protocol == "https" }) diff --git a/tools/register_command_test.go b/tools/register_command_test.go index fb29f21a..3fc4fd3b 100644 --- a/tools/register_command_test.go +++ b/tools/register_command_test.go @@ -4,6 +4,7 @@ import ( "bufio" "bytes" "context" + "encoding/json" "fmt" "io" "net" @@ -19,9 +20,10 @@ import ( "testing" "time" + aop "github.com/chainreactors/aiscan/aop" + toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/capability" "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/resources" "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" @@ -255,7 +257,7 @@ func TestNeutronSetProxyUpdatesDefault(t *testing.T) { type functionalResult struct { Stdout string Stderr string - Events []output.ToolDataEvent + Events []functionalEvent } type functionalCase struct { @@ -269,29 +271,66 @@ type functionalCase struct { type functionalRecorder struct { mu sync.Mutex - events []output.ToolDataEvent + events []functionalEvent } -func newFunctionalRecorder(bus *eventbus.Bus[output.ToolDataEvent]) *functionalRecorder { +type functionalEvent struct { + Tool, Kind, Target, CallID string + Data any +} + +func newFunctionalRecorder(bus *eventbus.Bus[*aop.Event]) *functionalRecorder { recorder := &functionalRecorder{} - bus.Subscribe(func(event output.ToolDataEvent) { + bus.Subscribe(func(event *aop.Event) { + if event == nil || event.GetExtension() == nil { + return + } + artifact := new(toolpb.Artifact) + if event.GetExtension().UnmarshalTo(artifact) != nil { + return + } + decoded := decodeFunctionalArtifact(artifact) recorder.mu.Lock() - recorder.events = append(recorder.events, event) + recorder.events = append(recorder.events, functionalEvent{ + Tool: artifact.Tool, Kind: artifact.Kind, Target: artifact.Target, CallID: artifact.CallId, Data: decoded, + }) recorder.mu.Unlock() }) return recorder } +func decodeFunctionalArtifact(artifact *toolpb.Artifact) any { + if artifact == nil { + return nil + } + var value any + switch artifact.Tool { + case "gogo": + value = new(parsers.GOGOResult) + case "spray": + value = new(parsers.SprayResult) + default: + value = new(any) + } + if json.Unmarshal(artifact.Data, value) != nil { + return nil + } + if holder, ok := value.(*any); ok { + return *holder + } + return value +} + func (r *functionalRecorder) mark() int { r.mu.Lock() defer r.mu.Unlock() return len(r.events) } -func (r *functionalRecorder) since(mark int) []output.ToolDataEvent { +func (r *functionalRecorder) since(mark int) []functionalEvent { r.mu.Lock() defer r.mu.Unlock() - return append([]output.ToolDataEvent(nil), r.events[mark:]...) + return append([]functionalEvent(nil), r.events[mark:]...) } func runFunctionalCases(t *testing.T, registry *commands.CommandRegistry, recorder *functionalRecorder, cases []functionalCase) { @@ -358,7 +397,7 @@ func requireOutputContains(t *testing.T, result functionalResult, values ...stri } } -func requireEvent(t *testing.T, result functionalResult, tool, kind string, match func(any) bool) output.ToolDataEvent { +func requireEvent(t *testing.T, result functionalResult, tool, kind string, match func(any) bool) functionalEvent { t.Helper() for _, event := range result.Events { if event.Tool == tool && event.Kind == kind && (match == nil || match(event.Data)) { @@ -366,10 +405,10 @@ func requireEvent(t *testing.T, result functionalResult, tool, kind string, matc } } t.Fatalf("missing event tool=%s kind=%s in %s", tool, kind, formatFunctionalEvents(result.Events)) - return output.ToolDataEvent{} + return functionalEvent{} } -func formatFunctionalEvents(events []output.ToolDataEvent) string { +func formatFunctionalEvents(events []functionalEvent) string { var b strings.Builder for _, event := range events { fmt.Fprintf(&b, "{%s %s %s %T} ", event.Tool, event.Kind, event.Target, event.Data) @@ -414,12 +453,12 @@ func TestScannerFunctionalRegression(t *testing.T) { defer engineSet.Close() workDir := t.TempDir() - bus := eventbus.New[output.ToolDataEvent]() + bus := eventbus.New[*aop.Event]() recorder := newFunctionalRecorder(bus) registry := commands.NewRegistry() deps := &commands.Deps{ WorkDir: workDir, - DataBus: bus, + Events: bus, Logger: telemetry.NopLogger(), } commands.Provide(deps, engine.SetKey, engineSet) @@ -459,7 +498,7 @@ http: Args: []string{"-i", host, "-p", port, "-v", "-o", "jl", "-t", "20"}, Check: func(t *testing.T, result functionalResult) { requireOutputContains(t, result, `"port":"`+port+`"`, "nginx") - requireEvent(t, result, "gogo", output.ToolDataService, func(data any) bool { + requireEvent(t, result, "gogo", toolpb.ArtifactKindService, func(data any) bool { item, ok := data.(*parsers.GOGOResult) if !ok || item == nil || item.Port != port { return false @@ -481,7 +520,7 @@ http: Args: []string{"-u", httpServer.URL, "--finger", "-j", "--limit", "5"}, Check: func(t *testing.T, result functionalResult) { requireOutputContains(t, result, httpServer.URL, "nginx") - requireEvent(t, result, "spray", output.ToolDataWeb, func(data any) bool { + requireEvent(t, result, "spray", toolpb.ArtifactKindWeb, func(data any) bool { item, ok := data.(*parsers.SprayResult) if !ok || item == nil || item.Status != http.StatusOK { return false @@ -496,7 +535,7 @@ http: Args: []string{"-u", tlsServer.URL, "-j", "--limit", "1"}, Check: func(t *testing.T, result functionalResult) { requireOutputContains(t, result, `"url":"`+tlsServer.URL+`"`, `"status":200`) - requireEvent(t, result, "spray", output.ToolDataWeb, func(data any) bool { + requireEvent(t, result, "spray", toolpb.ArtifactKindWeb, func(data any) bool { item, ok := data.(*parsers.SprayResult) return ok && item != nil && item.Status == http.StatusOK && strings.HasPrefix(item.UrlString, "https://") }) @@ -533,7 +572,7 @@ http: Args: []string{"-i", httpServer.URL, "-t", templateFile, "--tags", "regression", "-s", "high", "-j"}, Check: func(t *testing.T, result functionalResult) { requireOutputContains(t, result, `"matched":true`, `"template":"regression-marker"`) - requireEvent(t, result, "neutron", output.ToolDataVuln, nil) + requireEvent(t, result, "neutron", toolpb.ArtifactKindVuln, nil) }, }, { @@ -541,7 +580,7 @@ http: Args: []string{"-i", secretFile, "-j"}, Check: func(t *testing.T, result functionalResult) { requireOutputContains(t, result, "AKIAIOSFODNN7EXAMPLE") - requireEvent(t, result, "proton", output.ToolDataVuln, nil) + requireEvent(t, result, "proton", toolpb.ArtifactKindVuln, nil) }, }, { @@ -557,7 +596,7 @@ http: Timeout: 30 * time.Second, Check: func(t *testing.T, result functionalResult) { requireOutputContains(t, result, "[summary] completed", port) - requireEvent(t, result, "gogo", output.ToolDataService, nil) + requireEvent(t, result, "gogo", toolpb.ArtifactKindService, nil) }, }, } diff --git a/tools/scan/command.go b/tools/scan/command.go index adb7a515..d398eb4f 100644 --- a/tools/scan/command.go +++ b/tools/scan/command.go @@ -4,10 +4,9 @@ import ( "context" "fmt" "io" - "os" - "path/filepath" "github.com/chainreactors/aiscan/agent" + toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/telemetry" @@ -36,7 +35,6 @@ type flags struct { Trace bool `long:"trace" description:"Show internal scanner source and pipeline trace"` Debug bool `long:"debug" description:"Enable trace and underlying scanner debug logs"` JSON bool `short:"j" long:"json" description:"Output raw gogo and spray results as JSON Lines"` - OutputFile string `short:"f" long:"file" description:"Write raw scanner records as JSON Lines"` NoColor bool `long:"no-color" description:"Disable ANSI colors in terminal output"` Ports string `long:"ports" description:"Ports for gogo scanning; defaults to all in quick and - in full"` Threads int // derived from Thread; not a CLI flag @@ -186,31 +184,36 @@ func (c *Command) execute(ctx context.Context, args []string, stream io.Writer) } result := coll.StructuredResult() c.emitStructuredData(ctx, result) - if flags.OutputFile != "" { - raw, outputErr := coll.JSONLines() - if outputErr != nil { - c.Logger.Errorf("scan output file: %s", outputErr) - } else if err := writeOutputFile(flags.OutputFile, raw); err != nil { - c.Logger.Errorf("%s", err.Error()) - } - } return out, result, nil } func (c *Command) emitStructuredData(ctx context.Context, result *output.ScanResult) { - if result == nil || c.DataBus == nil { + if result == nil || c.Events == nil { return } for _, service := range result.GOGO { if service != nil { - c.EmitDataCtx(ctx, "gogo", output.ToolDataService, service.GetTarget(), service) + c.EmitArtifactCtx(ctx, "gogo", toolpb.ArtifactKindService, service.GetTarget(), service) } } for _, probe := range result.Spray { if probe != nil { - c.EmitDataCtx(ctx, "spray", output.ToolDataWeb, probe.UrlString, probe) + c.EmitArtifactCtx(ctx, "spray", toolpb.ArtifactKindWeb, probe.UrlString, probe) + } + } + for i := range result.Loots { + loot := result.Loots[i] + kind := loot.Kind + if kind == "" { + kind = toolpb.ArtifactKindVuln } + c.EmitArtifactCtx(ctx, "scan", kind, loot.Target, &loot) } + for i := range result.Errors { + scanErr := result.Errors[i] + c.EmitArtifactCtx(ctx, "scan", toolpb.ArtifactKindError, scanErr.Source, &scanErr) + } + c.EmitArtifactCtx(ctx, "scan", toolpb.ArtifactKindSummary, "", &result.Summary) } var scanFileFlags = map[string]bool{ @@ -222,28 +225,3 @@ var scanFileFlags = map[string]bool{ func (c *Command) resolveRelativePaths(args []string) []string { return toolargs.ResolveRelativePaths(args, scanFileFlags, c.WorkDir) } - -func writeOutputFile(path, content string) error { - path = filepath.Clean(path) - if dir := filepath.Dir(path); dir != "." && dir != "" { - if err := os.MkdirAll(dir, 0755); err != nil { - return fmt.Errorf("scan output file: create directory: %w", err) - } - } - f, err := os.Create(path) - if err != nil { - return fmt.Errorf("scan output file: %w", err) - } - if _, err := io.WriteString(f, content); err != nil { - _ = f.Close() - return fmt.Errorf("scan output file: write: %w", err) - } - if err := f.Sync(); err != nil { - _ = f.Close() - return fmt.Errorf("scan output file: sync: %w", err) - } - if err := f.Close(); err != nil { - return fmt.Errorf("scan output file: close: %w", err) - } - return nil -} diff --git a/tools/scan/command_test.go b/tools/scan/command_test.go index 89b5fa74..5ad4a9b2 100644 --- a/tools/scan/command_test.go +++ b/tools/scan/command_test.go @@ -14,9 +14,12 @@ import ( "testing" "time" + aop "github.com/chainreactors/aiscan/aop" + toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/telemetry" + coretool "github.com/chainreactors/aiscan/core/tool" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/tools/scan/engine" "github.com/chainreactors/aiscan/tools/scan/pipeline" @@ -1438,43 +1441,10 @@ func TestStructuredResultKeepsScannerValuesInsideCollector(t *testing.T) { } } -func TestScanOutputFileContainsRawRecordsWithoutChangingStdout(t *testing.T) { - sprayEng, _ := spray.NewEngine(nil) - cmd := New(&engine.Set{Spray: sprayEng}) - file := filepath.Join(t.TempDir(), "scan.txt") - var stdout bytes.Buffer - details, err := cmd.Run(context.Background(), &commands.Execution{Args: []string{"-i", "http://127.0.0.1:1", "--mode", "quick", "--timeout", "1", "-f", file}, Stdout: &stdout, Stderr: &stdout}) - if err != nil { - t.Fatalf("Run() error = %v", err) - } - // Structured scan records flow through the artifact stream; Run no longer - // returns a second result envelope. - if details != nil { - t.Fatalf("Run() returned unexpected details: %#v", details) - } - out := stdout.String() - data, err := os.ReadFile(file) - if err != nil { - t.Fatalf("read output file: %v", err) - } - fileOut := string(data) - if hasANSI(fileOut) { - t.Fatalf("file output contains ANSI: %q", fileOut) - } - if strings.Contains(fileOut, "[summary]") || strings.Contains(fileOut, "scan_start") || strings.Contains(fileOut, "scan_end") { - t.Fatalf("raw record file contains presentation data: %q", fileOut) - } - if !strings.Contains(output.StripANSI(out), "[summary] completed") { - t.Fatalf("stdout output missing summary: %q", out) - } - if strings.Contains(out, "[scan.web] ") { - t.Fatalf("stdout output should not repeat streamed events: %q", out) - } - if !strings.Contains(output.StripANSI(out), "http://127.0.0.1:1") { - t.Fatalf("stdout missing event line: %q", out) - } - if strings.Contains(output.StripANSI(out), "type=web") { - t.Fatalf("stdout contains key/value pollution: %q", out) +func TestScanCommandDoesNotExposeIndependentFileWriter(t *testing.T) { + usage := Usage() + if strings.Contains(usage, "--file") || strings.Contains(usage, "/file") || strings.Contains(usage, "Write raw scanner records") { + t.Fatalf("scan usage still exposes a side-channel file writer:\n%s", usage) } } @@ -1684,16 +1654,18 @@ func TestCleanupGogoTempFilesIgnoresMissingFile(t *testing.T) { } func TestEmitStructuredDataPublishesScannerFacts(t *testing.T) { - bus := eventbus.New[output.ToolDataEvent]() - cmd := New(&engine.Set{}, WithDataBus(bus)) + bus := eventbus.New[*aop.Event]() + cmd := New(&engine.Set{}, WithEvents(bus)) - var events []output.ToolDataEvent - unsub := bus.Subscribe(func(event output.ToolDataEvent) { + var events []*aop.Event + unsub := bus.Subscribe(func(event *aop.Event) { events = append(events, event) }) defer unsub() - ctx := output.ContextWithCallID(context.Background(), "scan-call-1") + ctx := coretool.ContextWithInvocation(context.Background(), coretool.Invocation{ + CallID: "scan-call-1", SessionID: "scan-session", TurnID: "scan-turn", Emitter: "scan", + }) cmd.emitStructuredData(ctx, &output.ScanResult{ GOGO: []*parsers.GOGOResult{{Ip: "127.0.0.1", Port: "8080", Protocol: "http"}}, Spray: []*parsers.SprayResult{{ @@ -1701,18 +1673,19 @@ func TestEmitStructuredDataPublishesScannerFacts(t *testing.T) { }}, }) - if len(events) != 2 { - t.Fatalf("events = %d, want 2: %#v", len(events), events) - } - if events[0].Tool != "gogo" || events[0].Kind != output.ToolDataService { - t.Fatalf("service event = %#v", events[0]) + if len(events) != 3 { + t.Fatalf("events = %d, want 3: %#v", len(events), events) } - if events[1].Tool != "spray" || events[1].Kind != output.ToolDataWeb { - t.Fatalf("web event = %#v", events[1]) + wants := []struct{ tool, kind string }{ + {"gogo", toolpb.ArtifactKindService}, {"spray", toolpb.ArtifactKindWeb}, {"scan", toolpb.ArtifactKindSummary}, } - for _, event := range events { - if event.CallID != "scan-call-1" { - t.Fatalf("call id = %q, want scan-call-1", event.CallID) + for index, event := range events { + artifact := new(toolpb.Artifact) + if event.GetExtension() == nil || event.GetExtension().UnmarshalTo(artifact) != nil { + t.Fatalf("artifact event = %#v", event) + } + if artifact.Tool != wants[index].tool || artifact.Kind != wants[index].kind || artifact.CallId != "scan-call-1" { + t.Fatalf("artifact = %#v, want %#v", artifact, wants[index]) } } } diff --git a/tools/scan/options.go b/tools/scan/options.go index cc82d83d..e9870e1c 100644 --- a/tools/scan/options.go +++ b/tools/scan/options.go @@ -4,8 +4,7 @@ import ( "context" "github.com/chainreactors/aiscan/agent" - "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/core/output" + aop "github.com/chainreactors/aiscan/aop" "github.com/chainreactors/aiscan/core/telemetry" ) @@ -21,8 +20,8 @@ func WithProxy(proxy string) Option { return func(c *Command) { c.Proxy = proxy } } -func WithDataBus(bus *eventbus.Bus[output.ToolDataEvent]) Option { - return func(c *Command) { c.DataBus = bus } +func WithEvents(events aop.EventEmitter) Option { + return func(c *Command) { c.Events = events } } func WithLogger(logger telemetry.Logger) Option { diff --git a/tools/spray/register.go b/tools/spray/register.go index 9b9f6470..8edf577b 100644 --- a/tools/spray/register.go +++ b/tools/spray/register.go @@ -21,7 +21,7 @@ func init() { d.Skip("spray", deps.Name(engine.SetKey)+".Spray") return } - impl := New(es.Spray).WithLogger(d.GetLogger()).WithProxy(d.ScannerProxy).WithDataBus(d.DataBus) + impl := New(es.Spray).WithLogger(d.GetLogger()).WithProxy(d.ScannerProxy).WithEvents(d.Events) reg.Register(commands.Command{ Name: impl.Name(), Usage: impl.Usage(), QuickReference: impl.QuickReference(), DescriptionPath: "aiscan://skills/aiscan/okf/easm/spray.md", diff --git a/tools/spray/spray.go b/tools/spray/spray.go index 142f016e..85e4910b 100644 --- a/tools/spray/spray.go +++ b/tools/spray/spray.go @@ -7,8 +7,8 @@ import ( "io" "strings" - "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/core/output" + aop "github.com/chainreactors/aiscan/aop" + toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/tools/toolargs" @@ -38,8 +38,8 @@ func (c *Command) WithProxy(proxy string) *Command { return c } -func (c *Command) WithDataBus(bus *eventbus.Bus[output.ToolDataEvent]) *Command { - c.DataBus = bus +func (c *Command) WithEvents(events aop.EventEmitter) *Command { + c.Events = events return c } @@ -110,7 +110,7 @@ func (c *Command) Run(ctx context.Context, execution *commands.Execution) (_ any return nil }, OnResult: func(r *parsers.SprayResult) { - c.EmitDataCtx(ctx, "spray", output.ToolDataWeb, r.UrlString, r) + c.EmitArtifactCtx(ctx, "spray", toolpb.ArtifactKindWeb, r.UrlString, r) if debug { // spray core prints results itself when not quiet return diff --git a/tools/toolargs/base.go b/tools/toolargs/base.go index 94a1e5a1..2e561635 100644 --- a/tools/toolargs/base.go +++ b/tools/toolargs/base.go @@ -2,18 +2,22 @@ package toolargs import ( "context" + "encoding/json" "time" - "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/core/output" + aop "github.com/chainreactors/aiscan/aop" + toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/telemetry" + coretool "github.com/chainreactors/aiscan/core/tool" + "google.golang.org/protobuf/types/known/anypb" + "google.golang.org/protobuf/types/known/timestamppb" ) type Base struct { Logger telemetry.Logger Proxy string WorkDir string - DataBus *eventbus.Bus[output.ToolDataEvent] + Events aop.EventEmitter } func (b *Base) SetWorkDir(dir string) { b.WorkDir = dir } @@ -27,20 +31,31 @@ func (b *Base) InitLogger(logger telemetry.Logger) { } } -func (b *Base) EmitData(tool, kind, target string, data any) { - b.EmitDataCtx(context.Background(), tool, kind, target, data) -} - -func (b *Base) EmitDataCtx(ctx context.Context, tool, kind, target string, data any) { - if b.DataBus == nil { +func (b *Base) EmitArtifactCtx(ctx context.Context, tool, kind, target string, data any) { + if b.Events == nil || data == nil { + return + } + raw, err := json.Marshal(data) + if err != nil { + b.Logger.Warnf("marshal %s artifact: %s", tool, err) return } - b.DataBus.Emit(output.ToolDataEvent{ - Tool: tool, - Kind: kind, - Target: target, - Data: data, - CallID: output.CallIDFromContext(ctx), - Timestamp: time.Now(), + invocation := coretool.InvocationFromContext(ctx) + artifact := &toolpb.Artifact{ + Tool: tool, Kind: kind, Target: target, Data: raw, + MediaType: aop.JSONMediaType, Timestamp: timestamppb.New(time.Now()), CallId: invocation.CallID, + } + extension, err := anypb.New(artifact) + if err != nil { + b.Logger.Warnf("encode %s artifact: %s", tool, err) + return + } + emitter := invocation.Emitter + if emitter == "" { + emitter = tool + } + b.Events.Emit(&aop.Event{ + SessionId: invocation.SessionID, TurnId: invocation.TurnID, Emitter: emitter, + Payload: &aop.Event_Extension{Extension: extension}, }) } diff --git a/tools/zombie/register.go b/tools/zombie/register.go index df76a389..c7456ab3 100644 --- a/tools/zombie/register.go +++ b/tools/zombie/register.go @@ -21,7 +21,7 @@ func init() { d.Skip("zombie", deps.Name(engine.SetKey)+".Zombie") return } - impl := New(es.Zombie).WithLogger(d.GetLogger()).WithProxy(d.ScannerProxy).WithDataBus(d.DataBus) + impl := New(es.Zombie).WithLogger(d.GetLogger()).WithProxy(d.ScannerProxy).WithEvents(d.Events) reg.Register(commands.Command{ Name: impl.Name(), Usage: impl.Usage(), DescriptionPath: "aiscan://skills/aiscan/okf/easm/zombie.md", diff --git a/tools/zombie/zombie.go b/tools/zombie/zombie.go index 00119323..e1519375 100644 --- a/tools/zombie/zombie.go +++ b/tools/zombie/zombie.go @@ -6,8 +6,7 @@ import ( "fmt" "os" - "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/core/output" + aop "github.com/chainreactors/aiscan/aop" "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/tools/toolargs" @@ -36,8 +35,8 @@ func (c *Command) WithProxy(proxy string) *Command { return c } -func (c *Command) WithDataBus(bus *eventbus.Bus[output.ToolDataEvent]) *Command { - c.DataBus = bus +func (c *Command) WithEvents(events aop.EventEmitter) *Command { + c.Events = events return c } diff --git a/web/frontend/cyber-ui b/web/frontend/cyber-ui index f476545d..de41e92b 160000 --- a/web/frontend/cyber-ui +++ b/web/frontend/cyber-ui @@ -1 +1 @@ -Subproject commit f476545d660ba61fefaa26190ce0fffb5ed04524 +Subproject commit de41e92bd43b554dab39be7de9e6ac7c36289418