diff --git a/crates/cli/src/banner.rs b/crates/cli/src/banner.rs index 180159cd..978bb3e0 100644 --- a/crates/cli/src/banner.rs +++ b/crates/cli/src/banner.rs @@ -51,6 +51,19 @@ const BORDER_BR: char = '╯'; const BORDER_H: char = '─'; const BORDER_V: char = '│'; +#[derive(Clone, Copy)] +struct DockTagSpan { + row: usize, + start: usize, + end: usize, +} + +enum CellStyle { + DockTag, + Figlet, + Plain, +} + fn supports_banner() -> bool { if !std::io::stdout().is_terminal() { return false; @@ -93,45 +106,10 @@ fn render_frame_inner(color: bool, docked: bool) -> String { let mut out = String::with_capacity(BANNER_LINES.iter().map(|l| l.len() + 64).sum()); out.push('\n'); - // Build a 2D grid: empty top rail, the 6 figlet rows, empty bottom rail. Each cell is a - // single char (we treat Unicode block chars as 1 display column wide, which is true for the - // glyphs the figlet uses). - let mut grid: Vec> = Vec::with_capacity(TOTAL_ROWS); let dock_tag = format!(" v{}", env!("CARGO_PKG_VERSION")); - let dock_width_needed = COL_END + dock_tag.chars().count() + 2; - let max_width = BANNER_LINES - .iter() - .map(|l| l.chars().count()) - .max() - .unwrap_or(0) - .max(dock_width_needed); - - // Top rail (empty). - grid.push(vec![' '; max_width]); - // 6 figlet rows, padded to max_width. - for line in BANNER_LINES { - let mut row: Vec = line.chars().collect(); - while row.len() < max_width { - row.push(' '); - } - grid.push(row); - } - // Bottom rail (empty). - grid.push(vec![' '; max_width]); - - // Overlay the docked version tag at bottom-right: just "vX.Y.Z" in dim green. No dot — the - // version reads as a quiet label below "Flow", letting the brand mark stand on its own. - let dock_col_start = COL_END; - let dock_col_end = dock_col_start + dock_tag.chars().count(); - if docked { - let dock_row = BOTTOM_RAIL; - for (i, ch) in dock_tag.chars().enumerate() { - let c = dock_col_start + i; - if dock_row < grid.len() && c < grid[dock_row].len() { - grid[dock_row][c] = ch; - } - } - } + let max_width = frame_width(&dock_tag); + let mut grid = build_grid(max_width); + let dock_tag_span = docked.then(|| overlay_dock_tag(&mut grid, &dock_tag)); // Top border row. push_border_line(&mut out, BORDER_TL, BORDER_TR, max_width, color); @@ -139,45 +117,7 @@ fn render_frame_inner(color: bool, docked: bool) -> String { // Emit the grid with appropriate coloring per cell. Each grid row is wrapped with a // vertical border on the left and right, painted in NVIDIA green. for (row_idx, row) in grid.iter().enumerate() { - if color { - out.push_str(NVIDIA_GREEN); - out.push(BORDER_V); - out.push_str(RESET); - } else { - out.push(BORDER_V); - } - for (col_idx, ch) in row.iter().enumerate() { - let in_dock_tag = docked - && row_idx == BOTTOM_RAIL - && col_idx >= dock_col_start - && col_idx < dock_col_end; - if in_dock_tag && *ch != ' ' { - if color { - out.push_str(DOCK_TAG); - out.push(*ch); - out.push_str(RESET); - } else { - out.push(*ch); - } - } else if is_figlet_glyph(*ch) { - if color { - out.push_str(NVIDIA_GREEN); - out.push(*ch); - out.push_str(RESET); - } else { - out.push(*ch); - } - } else { - out.push(*ch); - } - } - if color { - out.push_str(NVIDIA_GREEN); - out.push(BORDER_V); - out.push_str(RESET); - } else { - out.push(BORDER_V); - } + push_grid_row(&mut out, row_idx, row, dock_tag_span, color); out.push('\n'); } @@ -187,6 +127,102 @@ fn render_frame_inner(color: bool, docked: bool) -> String { out } +fn frame_width(dock_tag: &str) -> usize { + let dock_width_needed = COL_END + dock_tag.chars().count() + 2; + BANNER_LINES + .iter() + .map(|l| l.chars().count()) + .max() + .unwrap_or(0) + .max(dock_width_needed) +} + +fn build_grid(width: usize) -> Vec> { + // Empty top rail, the 6 figlet rows, and an empty bottom rail. Each cell is a single char + // because the figlet's block and box glyphs render as one display column in target terminals. + let mut grid = Vec::with_capacity(TOTAL_ROWS); + grid.push(vec![' '; width]); + grid.extend(BANNER_LINES.iter().map(|line| padded_row(line, width))); + grid.push(vec![' '; width]); + grid +} + +fn padded_row(line: &str, width: usize) -> Vec { + let mut row: Vec = line.chars().collect(); + row.resize(width, ' '); + row +} + +fn overlay_dock_tag(grid: &mut [Vec], dock_tag: &str) -> DockTagSpan { + let span = DockTagSpan { + row: BOTTOM_RAIL, + start: COL_END, + end: COL_END + dock_tag.chars().count(), + }; + for (index, ch) in dock_tag.chars().enumerate() { + grid[span.row][span.start + index] = ch; + } + span +} + +fn push_grid_row( + out: &mut String, + row_idx: usize, + row: &[char], + dock_tag_span: Option, + color: bool, +) { + push_vertical_border(out, color); + for (col_idx, ch) in row.iter().copied().enumerate() { + push_cell( + out, + ch, + cell_style(ch, row_idx, col_idx, dock_tag_span), + color, + ); + } + push_vertical_border(out, color); +} + +fn push_vertical_border(out: &mut String, color: bool) { + push_styled_char(out, BORDER_V, Some(NVIDIA_GREEN), color); +} + +fn push_cell(out: &mut String, ch: char, style: CellStyle, color: bool) { + match style { + CellStyle::DockTag => push_styled_char(out, ch, Some(DOCK_TAG), color), + CellStyle::Figlet => push_styled_char(out, ch, Some(NVIDIA_GREEN), color), + CellStyle::Plain => out.push(ch), + } +} + +fn push_styled_char(out: &mut String, ch: char, style: Option<&str>, color: bool) { + if color && let Some(style) = style { + out.push_str(style); + out.push(ch); + out.push_str(RESET); + } else { + out.push(ch); + } +} + +fn cell_style( + ch: char, + row_idx: usize, + col_idx: usize, + dock_tag_span: Option, +) -> CellStyle { + if dock_tag_span.is_some_and(|span| { + row_idx == span.row && col_idx >= span.start && col_idx < span.end && ch != ' ' + }) { + CellStyle::DockTag + } else if is_figlet_glyph(ch) { + CellStyle::Figlet + } else { + CellStyle::Plain + } +} + fn push_border_line(out: &mut String, left: char, right: char, inner_width: usize, color: bool) { if color { out.push_str(NVIDIA_GREEN); diff --git a/crates/cli/src/setup.rs b/crates/cli/src/setup.rs index f1c94be6..8a9183b1 100644 --- a/crates/cli/src/setup.rs +++ b/crates/cli/src/setup.rs @@ -116,60 +116,88 @@ pub(crate) fn detect_installed_agents_in(path_var: Option<&std::ffi::OsStr>) -> pub(crate) fn build_config(answers: &SetupAnswers) -> DocumentMut { let mut doc = DocumentMut::new(); - // Build the exporter table once so selecting multiple backends produces nested per-exporter - // sections, not separate legacy observability/export blocks. - let want_atif = answers.backends.contains(&ObservabilityBackend::Atif); - let want_atof = answers.backends.contains(&ObservabilityBackend::Atof); - let want_openinference = answers - .backends - .contains(&ObservabilityBackend::OpenInference) - && answers.openinference_endpoint.is_some(); - if want_atif || want_atof || want_openinference { - let mut exporters = Table::new(); - if want_atif { - let mut atif = Table::new(); - atif["dir"] = value("./atif"); - exporters.insert("atif", Item::Table(atif)); - } - if want_atof { - let mut atof = Table::new(); - atof["dir"] = value("./atof"); - atof["mode"] = value("append"); - atof["filename_template"] = value("{session_id}.jsonl"); - exporters.insert("atof", Item::Table(atof)); - } - if let Some(endpoint) = answers.openinference_endpoint.as_deref() { - let mut openinference = Table::new(); - openinference["endpoint"] = value(endpoint); - exporters.insert("openinference", Item::Table(openinference)); - } + if let Some(exporters) = build_exporters_table(answers) { doc["exporters"] = Item::Table(exporters); } - if !answers.agents.is_empty() { - let mut agents_table = Table::new(); - for agent in &answers.agents { - let (key, command) = match agent { - CodingAgent::ClaudeCode => ("claude", "claude"), - CodingAgent::Codex => ("codex", "codex"), - CodingAgent::Cursor => ("cursor", "cursor-agent"), - CodingAgent::Hermes => ("hermes", "hermes"), - }; - let mut agent_table = Table::new(); - agent_table["command"] = value(command); - if matches!(agent, CodingAgent::Hermes) - && let Some(path) = answers.hermes_hooks_path.as_deref() - { - agent_table["hooks_path"] = value(path.display().to_string()); - } - agents_table.insert(key, Item::Table(agent_table)); - } + if let Some(agents_table) = build_agents_table(answers) { doc["agents"] = Item::Table(agents_table); } doc } +fn build_exporters_table(answers: &SetupAnswers) -> Option { + if !has_selected_exporter(answers) { + return None; + } + + let mut exporters = Table::new(); + if answers.backends.contains(&ObservabilityBackend::Atif) { + insert_atif_exporter(&mut exporters); + } + if answers.backends.contains(&ObservabilityBackend::Atof) { + insert_atof_exporter(&mut exporters); + } + if answers + .backends + .contains(&ObservabilityBackend::OpenInference) + && let Some(endpoint) = answers.openinference_endpoint.as_deref() + { + insert_openinference_exporter(&mut exporters, endpoint); + } + Some(exporters) +} + +fn has_selected_exporter(answers: &SetupAnswers) -> bool { + answers.backends.contains(&ObservabilityBackend::Atif) + || answers.backends.contains(&ObservabilityBackend::Atof) + || (answers + .backends + .contains(&ObservabilityBackend::OpenInference) + && answers.openinference_endpoint.is_some()) +} + +fn insert_atif_exporter(exporters: &mut Table) { + let mut atif = Table::new(); + atif["dir"] = value("./atif"); + exporters.insert("atif", Item::Table(atif)); +} + +fn insert_atof_exporter(exporters: &mut Table) { + let mut atof = Table::new(); + atof["dir"] = value("./atof"); + atof["mode"] = value("append"); + atof["filename_template"] = value("{session_id}.jsonl"); + exporters.insert("atof", Item::Table(atof)); +} + +fn insert_openinference_exporter(exporters: &mut Table, endpoint: &str) { + let mut openinference = Table::new(); + openinference["endpoint"] = value(endpoint); + exporters.insert("openinference", Item::Table(openinference)); +} + +fn build_agents_table(answers: &SetupAnswers) -> Option
{ + if answers.agents.is_empty() { + return None; + } + + let mut agents_table = Table::new(); + for agent in &answers.agents { + let (key, command) = agent_key_and_command(*agent); + let mut agent_table = Table::new(); + agent_table["command"] = value(command); + if matches!(agent, CodingAgent::Hermes) + && let Some(path) = answers.hermes_hooks_path.as_deref() + { + agent_table["hooks_path"] = value(path.display().to_string()); + } + agents_table.insert(key, Item::Table(agent_table)); + } + Some(agents_table) +} + /// Writes the setup's TOML document to the scope-appropriate path(s). /// /// When `merge_scope` is `Some(agent)`, an existing `config.toml` at the target path is parsed diff --git a/crates/cli/tests/coverage/setup_tests.rs b/crates/cli/tests/coverage/setup_tests.rs index 403702c3..28b16305 100644 --- a/crates/cli/tests/coverage/setup_tests.rs +++ b/crates/cli/tests/coverage/setup_tests.rs @@ -114,6 +114,23 @@ fn build_config_emits_exporters_section_when_openinference_selected() { assert!(!rendered.contains("[observability]")); } +#[test] +fn build_config_ignores_openinference_endpoint_when_backend_not_selected() { + let answers = SetupAnswers { + scope: ConfigScope::Project, + agents: vec![], + backends: vec![ObservabilityBackend::Atif], + openinference_endpoint: Some("http://localhost:6006/v1/traces".into()), + hermes_hooks_path: None, + }; + + let rendered = build_config(&answers).to_string(); + + assert!(rendered.contains("[exporters.atif]")); + assert!(!rendered.contains("[exporters.openinference]")); + assert!(!rendered.contains("http://localhost:6006/v1/traces")); +} + #[test] fn build_config_emits_atof_write_options_when_atof_selected() { let answers = SetupAnswers { diff --git a/go/nemo_flow/atof_test.go b/go/nemo_flow/atof_test.go index cb879ab4..e8cfef4d 100644 --- a/go/nemo_flow/atof_test.go +++ b/go/nemo_flow/atof_test.go @@ -12,6 +12,8 @@ import ( "time" ) +const eventsJSONLFilename = "events.jsonl" + func TestNewAtofExporterConfigDefaults(t *testing.T) { config := NewAtofExporterConfig() @@ -31,47 +33,25 @@ func TestAtofExporterLifecycleWritesRawJSONL(t *testing.T) { exporter, err := NewAtofExporter(AtofExporterConfig{ OutputDirectory: dir, Mode: AtofExporterModeOverwrite, - Filename: "events.jsonl", + Filename: eventsJSONLFilename, }) - if err != nil { - t.Fatalf("NewAtofExporter failed: %v", err) - } + requireNoError(t, err, "NewAtofExporter failed") defer exporter.Close() path, err := exporter.Path() - if err != nil { - t.Fatalf("Path failed: %v", err) - } - if filepath.Base(path) != "events.jsonl" { - t.Fatalf("expected events.jsonl path, got %q", path) - } + requireNoError(t, err, "Path failed") + requireEqual(t, filepath.Base(path), eventsJSONLFilename, "expected %s path", eventsJSONLFilename) name := "go_atof_" + time.Now().Format("150405.000000") - if err := exporter.Register(name); err != nil { - t.Fatalf("Register failed: %v", err) - } + requireNoError(t, exporter.Register(name), "Register failed") handle, err := PushScope("atof_scope", ScopeTypeAgent, WithInput(json.RawMessage(`{"scope":true}`))) - if err != nil { - t.Fatalf("PushScope failed: %v", err) - } - if err := EmitEvent("atof_mark", WithEventParent(handle), WithEventData(json.RawMessage(`{"step":1}`))); err != nil { - t.Fatalf("EmitEvent failed: %v", err) - } - if err := PopScope(handle, WithOutput(json.RawMessage(`{"done":true}`))); err != nil { - t.Fatalf("PopScope failed: %v", err) - } - if err := exporter.Deregister(name); err != nil { - t.Fatalf("Deregister failed: %v", err) - } - if err := exporter.Deregister(name); err != nil { - t.Fatalf("repeated Deregister should be safe, got: %v", err) - } - if err := exporter.ForceFlush(); err != nil { - t.Fatalf("ForceFlush failed: %v", err) - } - if err := exporter.Shutdown(); err != nil { - t.Fatalf("Shutdown failed: %v", err) - } + requireNoError(t, err, "PushScope failed") + requireNoError(t, EmitEvent("atof_mark", WithEventParent(handle), WithEventData(json.RawMessage(`{"step":1}`))), "EmitEvent failed") + requireNoError(t, PopScope(handle, WithOutput(json.RawMessage(`{"done":true}`))), "PopScope failed") + requireNoError(t, exporter.Deregister(name), "Deregister failed") + requireNoError(t, exporter.Deregister(name), "repeated Deregister should be safe") + requireNoError(t, exporter.ForceFlush(), "ForceFlush failed") + requireNoError(t, exporter.Shutdown(), "Shutdown failed") records := readAtofRecords(t, path) if len(records) != 3 { @@ -90,14 +70,14 @@ func TestAtofExporterLifecycleWritesRawJSONL(t *testing.T) { func TestAtofExporterAppendAndOverwriteModes(t *testing.T) { dir := t.TempDir() - path := filepath.Join(dir, "events.jsonl") + path := filepath.Join(dir, eventsJSONLFilename) if err := os.WriteFile(path, []byte("{\"existing\":true}\n"), 0o600); err != nil { t.Fatalf("write seed file: %v", err) } appendExporter, err := NewAtofExporter(AtofExporterConfig{ OutputDirectory: dir, - Filename: "events.jsonl", + Filename: eventsJSONLFilename, }) if err != nil { t.Fatalf("append NewAtofExporter failed: %v", err) @@ -113,7 +93,7 @@ func TestAtofExporterAppendAndOverwriteModes(t *testing.T) { overwriteExporter, err := NewAtofExporter(AtofExporterConfig{ OutputDirectory: dir, Mode: AtofExporterModeOverwrite, - Filename: "events.jsonl", + Filename: eventsJSONLFilename, }) if err != nil { t.Fatalf("overwrite NewAtofExporter failed: %v", err) @@ -153,3 +133,17 @@ func mustReadFile(t *testing.T, path string) []byte { } return content } + +func requireNoError(t *testing.T, err error, message string) { + t.Helper() + if err != nil { + t.Fatalf("%s: %v", message, err) + } +} + +func requireEqual[T comparable](t *testing.T, got T, want T, message string, args ...any) { + t.Helper() + if got != want { + t.Fatalf(message+", got %q", append(args, got)...) + } +} diff --git a/go/nemo_flow/observability_plugin_test.go b/go/nemo_flow/observability_plugin_test.go index 6e28aae4..34222fcb 100644 --- a/go/nemo_flow/observability_plugin_test.go +++ b/go/nemo_flow/observability_plugin_test.go @@ -11,6 +11,15 @@ import ( "testing" ) +const ( + ClearPluginConfigurationFailed = "ClearPluginConfiguration failed" + InitializePluginsFailed = "InitializePlugins failed" + TrajectoryFilenamePrefix = "trajectory-" + FirstAgentName = "go-first-agent" + NestedAgentName = "go-nested-agent" + SecondAgentName = "go-second-agent" +) + func TestObservabilityConfigHelpers(t *testing.T) { config := NewObservabilityConfig() if config.Version != 1 { @@ -41,14 +50,17 @@ func TestObservabilityConfigHelpers(t *testing.T) { func TestObservabilityPluginAtofAndAtifFiles(t *testing.T) { if err := ClearPluginConfiguration(); err != nil { - t.Fatalf("ClearPluginConfiguration failed: %v", err) + t.Fatalf("%s: %v", ClearPluginConfigurationFailed, err) } + t.Cleanup(func() { + requireNoError(t, ClearPluginConfiguration(), ClearPluginConfigurationFailed) + }) dir := t.TempDir() config := NewObservabilityConfig() atof := NewObservabilityAtofConfig() atof.Enabled = true atof.OutputDirectory = dir - atof.Filename = "events.jsonl" + atof.Filename = eventsJSONLFilename atof.Mode = "overwrite" config.Atof = &atof atif := NewObservabilityAtifConfig() @@ -59,7 +71,7 @@ func TestObservabilityPluginAtofAndAtifFiles(t *testing.T) { atif.ToolDefinitions = []map[string]any{{"name": "search"}} atif.Extra = map[string]any{"binding": "go"} atif.OutputDirectory = dir - atif.FilenameTemplate = "trajectory-{session_id}.json" + atif.FilenameTemplate = TrajectoryFilenamePrefix + "{session_id}.json" config.Atif = &atif if report, err := ValidatePluginConfig(PluginConfig{Version: 1, Components: []PluginComponentSpec{ObservabilityComponent(config)}}); err != nil { @@ -68,7 +80,7 @@ func TestObservabilityPluginAtofAndAtifFiles(t *testing.T) { t.Fatalf("unexpected diagnostics: %#v", report.Diagnostics) } if _, err := InitializePlugins(PluginConfig{Version: 1, Components: []PluginComponentSpec{ObservabilityComponent(config)}}); err != nil { - t.Fatalf("InitializePlugins failed: %v", err) + t.Fatalf("%s: %v", InitializePluginsFailed, err) } handle, err := PushScope("go-observability-agent", ScopeTypeAgent, WithInput(json.RawMessage(`{"agent":true}`))) @@ -82,15 +94,15 @@ func TestObservabilityPluginAtofAndAtifFiles(t *testing.T) { t.Fatalf("PopScope failed: %v", err) } if err := ClearPluginConfiguration(); err != nil { - t.Fatalf("ClearPluginConfiguration failed: %v", err) + t.Fatalf("%s: %v", ClearPluginConfigurationFailed, err) } - jsonl := string(mustReadFile(t, filepath.Join(dir, "events.jsonl"))) + jsonl := string(mustReadFile(t, filepath.Join(dir, eventsJSONLFilename))) if got := strings.Count(strings.TrimSpace(jsonl), "\n") + 1; got != 3 { t.Fatalf("expected 3 JSONL records, got %d: %s", got, jsonl) } - trajectoryPath := filepath.Join(dir, "trajectory-"+handle.UUID()+".json") + trajectoryPath := TrajectoryFilePath(dir, handle) var trajectory map[string]any if err := json.Unmarshal(mustReadFile(t, trajectoryPath), &trajectory); err != nil { t.Fatalf("failed to read trajectory: %v", err) @@ -105,77 +117,36 @@ func TestObservabilityPluginAtofAndAtifFiles(t *testing.T) { } func TestObservabilityPluginAtifSplitsMultipleTopLevelAgents(t *testing.T) { - if err := ClearPluginConfiguration(); err != nil { - t.Fatalf("ClearPluginConfiguration failed: %v", err) - } - dir := t.TempDir() - config := NewObservabilityConfig() - atif := NewObservabilityAtifConfig() - atif.Enabled = true - atif.OutputDirectory = dir - atif.FilenameTemplate = "trajectory-{session_id}.json" - config.Atif = &atif + Dir := t.TempDir() + InitializeAtifPlugin(t, Dir) + First := EmitAgentStart(t, "first", FirstAgentName) + Nested := EmitAgentStart(t, "nested", NestedAgentName) + EmitAgentEnd(t, "nested", Nested) + EmitAgentEnd(t, "first", First) + Second := EmitAgentTrajectory(t, "second", SecondAgentName) + requireNoError(t, ClearPluginConfiguration(), ClearPluginConfigurationFailed) - if _, err := InitializePlugins(PluginConfig{Version: 1, Components: []PluginComponentSpec{ObservabilityComponent(config)}}); err != nil { - t.Fatalf("InitializePlugins failed: %v", err) - } - - first, err := PushScope("go-first-agent", ScopeTypeAgent, WithInput(json.RawMessage(`{"agent":"first"}`))) - if err != nil { - t.Fatalf("PushScope first failed: %v", err) - } - if err := EmitEvent("go-first-mark", WithEventParent(first), WithEventData(json.RawMessage(`{"agent":"first"}`))); err != nil { - t.Fatalf("EmitEvent first failed: %v", err) - } - nested, err := PushScope("go-nested-agent", ScopeTypeAgent, WithInput(json.RawMessage(`{"agent":"nested"}`))) - if err != nil { - t.Fatalf("PushScope nested failed: %v", err) - } - if err := EmitEvent("go-nested-mark", WithEventParent(nested), WithEventData(json.RawMessage(`{"agent":"nested"}`))); err != nil { - t.Fatalf("EmitEvent nested failed: %v", err) - } - if err := PopScope(nested, WithOutput(json.RawMessage(`{"done":true}`))); err != nil { - t.Fatalf("PopScope nested failed: %v", err) - } - if err := PopScope(first, WithOutput(json.RawMessage(`{"done":true}`))); err != nil { - t.Fatalf("PopScope first failed: %v", err) - } - - second, err := PushScope("go-second-agent", ScopeTypeAgent, WithInput(json.RawMessage(`{"agent":"second"}`))) - if err != nil { - t.Fatalf("PushScope second failed: %v", err) - } - if err := EmitEvent("go-second-mark", WithEventParent(second), WithEventData(json.RawMessage(`{"agent":"second"}`))); err != nil { - t.Fatalf("EmitEvent second failed: %v", err) - } - if err := PopScope(second, WithOutput(json.RawMessage(`{"done":true}`))); err != nil { - t.Fatalf("PopScope second failed: %v", err) - } - if err := ClearPluginConfiguration(); err != nil { - t.Fatalf("ClearPluginConfiguration failed: %v", err) - } - - files, err := filepath.Glob(filepath.Join(dir, "trajectory-*.json")) + Files, err := filepath.Glob(filepath.Join(Dir, TrajectoryFilenamePrefix+"*.json")) if err != nil { t.Fatalf("Glob failed: %v", err) } - if len(files) != 2 { - t.Fatalf("expected 2 ATIF trajectory files, got %d: %#v", len(files), files) + if len(Files) != 2 { + t.Fatalf("expected 2 ATIF trajectory files, got %d: %#v", len(Files), Files) } - firstPayload := string(mustReadFile(t, filepath.Join(dir, "trajectory-"+first.UUID()+".json"))) - secondPayload := string(mustReadFile(t, filepath.Join(dir, "trajectory-"+second.UUID()+".json"))) - if !strings.Contains(firstPayload, "go-first-agent") || !strings.Contains(firstPayload, "go-nested-agent") { - t.Fatalf("expected first trajectory to include first and nested agents: %s", firstPayload) + FirstPayload := string(mustReadFile(t, TrajectoryFilePath(Dir, First))) + SecondPayload := string(mustReadFile(t, TrajectoryFilePath(Dir, Second))) + if !strings.Contains(FirstPayload, FirstAgentName) || !strings.Contains(FirstPayload, NestedAgentName) { + t.Fatalf("expected first trajectory to include first and nested agents: %s", FirstPayload) } - if strings.Contains(firstPayload, "go-second-agent") { - t.Fatalf("first trajectory leaked second agent events: %s", firstPayload) + if strings.Contains(FirstPayload, SecondAgentName) { + t.Fatalf("first trajectory leaked second agent events: %s", FirstPayload) } - if !strings.Contains(secondPayload, "go-second-agent") { - t.Fatalf("expected second trajectory to include second agent: %s", secondPayload) + if !strings.Contains(SecondPayload, SecondAgentName) { + t.Fatalf("expected second trajectory to include second agent: %s", SecondPayload) } - if strings.Contains(secondPayload, "go-first-agent") || strings.Contains(secondPayload, "go-nested-agent") { - t.Fatalf("second trajectory leaked first trajectory events: %s", secondPayload) + if strings.Contains(SecondPayload, FirstAgentName) || strings.Contains(SecondPayload, NestedAgentName) { + t.Fatalf("second trajectory leaked first trajectory events: %s", SecondPayload) } } @@ -212,8 +183,11 @@ func TestObservabilityPluginListKindIsAutomatic(t *testing.T) { func TestObservabilityAtifOpenAgentFlushesOnClear(t *testing.T) { if err := ClearPluginConfiguration(); err != nil { - t.Fatalf("ClearPluginConfiguration failed: %v", err) + t.Fatalf("%s: %v", ClearPluginConfigurationFailed, err) } + t.Cleanup(func() { + requireNoError(t, ClearPluginConfiguration(), ClearPluginConfigurationFailed) + }) dir := t.TempDir() config := NewObservabilityConfig() atif := NewObservabilityAtifConfig() @@ -221,14 +195,14 @@ func TestObservabilityAtifOpenAgentFlushesOnClear(t *testing.T) { atif.OutputDirectory = dir config.Atif = &atif if _, err := InitializePlugins(PluginConfig{Version: 1, Components: []PluginComponentSpec{ObservabilityComponent(config)}}); err != nil { - t.Fatalf("InitializePlugins failed: %v", err) + t.Fatalf("%s: %v", InitializePluginsFailed, err) } handle, err := PushScope("go-open-agent", ScopeTypeAgent) if err != nil { t.Fatalf("PushScope failed: %v", err) } if err := ClearPluginConfiguration(); err != nil { - t.Fatalf("ClearPluginConfiguration failed: %v", err) + t.Fatalf("%s: %v", ClearPluginConfigurationFailed, err) } path := filepath.Join(dir, "nemo-flow-atif-"+handle.UUID()+".json") if _, err := os.Stat(path); err != nil { @@ -238,3 +212,45 @@ func TestObservabilityAtifOpenAgentFlushesOnClear(t *testing.T) { t.Fatalf("PopScope failed: %v", err) } } + +func InitializeAtifPlugin(t *testing.T, Dir string) { + t.Helper() + t.Cleanup(func() { + requireNoError(t, ClearPluginConfiguration(), ClearPluginConfigurationFailed) + }) + requireNoError(t, ClearPluginConfiguration(), ClearPluginConfigurationFailed) + + Config := NewObservabilityConfig() + Atif := NewObservabilityAtifConfig() + Atif.Enabled = true + Atif.OutputDirectory = Dir + Atif.FilenameTemplate = TrajectoryFilenamePrefix + "{session_id}.json" + Config.Atif = &Atif + + _, Err := InitializePlugins(PluginConfig{Version: 1, Components: []PluginComponentSpec{ObservabilityComponent(Config)}}) + requireNoError(t, Err, InitializePluginsFailed) +} + +func EmitAgentTrajectory(t *testing.T, Label string, Name string) *ScopeHandle { + t.Helper() + Handle := EmitAgentStart(t, Label, Name) + EmitAgentEnd(t, Label, Handle) + return Handle +} + +func EmitAgentStart(t *testing.T, Label string, Name string) *ScopeHandle { + t.Helper() + Handle, Err := PushScope(Name, ScopeTypeAgent, WithInput(json.RawMessage(`{"agent":"`+Label+`"}`))) + requireNoError(t, Err, "PushScope "+Label+" failed") + requireNoError(t, EmitEvent("go-"+Label+"-mark", WithEventParent(Handle), WithEventData(json.RawMessage(`{"agent":"`+Label+`"}`))), "EmitEvent "+Label+" failed") + return Handle +} + +func EmitAgentEnd(t *testing.T, Label string, Handle *ScopeHandle) { + t.Helper() + requireNoError(t, PopScope(Handle, WithOutput(json.RawMessage(`{"done":true}`))), "PopScope "+Label+" failed") +} + +func TrajectoryFilePath(Dir string, Handle *ScopeHandle) string { + return filepath.Join(Dir, TrajectoryFilenamePrefix+Handle.UUID()+".json") +}