diff --git a/apps/maple-agent/crates/UPSTREAM.md b/apps/maple-agent/crates/UPSTREAM.md index a6a776236..e6552e42f 100644 --- a/apps/maple-agent/crates/UPSTREAM.md +++ b/apps/maple-agent/crates/UPSTREAM.md @@ -5,10 +5,12 @@ product policy, storage formats, account rules and UI reconciliation belong in `maple-agent` and the app. A generic API still needs an explicit reason when it adds behavior absent from Pi. -The initial reference is Pi `v1.0.4` -(`7c10bd4337495ee613f2224843ecdf349b80d1df`). The initial port and its accepted -language/API differences are recorded in -[Maple #1104](https://github.com/MaplePrivacyLabs/Maple/pull/1104). +The crates follow Pi `v1.4.0` (`4a080252f06caf63a525d7873c5d94d3c23887e0`, +2026-10-22), the end of the latest full review below; start the next review +there. The initial port from Pi `v1.0.4` and its accepted language/API +differences are recorded in +[Maple #1104](https://github.com/MaplePrivacyLabs/Maple/pull/1104); each +review's PR records its adopt/skip decisions and any new differences. This file records the storage follow-up's adaptations; it does not claim a complete audit of every earlier difference. @@ -55,3 +57,14 @@ There is no automatic line-by-line update or requirement to adopt every release. (`71023aec1546d6d677a65669c5665c0927ce0dd0`): ordinary message persistence still does not emit it. This was a focused contract check, not a full upstream review or a Pi update. +- 2026-10-30: full review of Pi `v1.0.4`..`v1.4.0` (`packages/ai/src`, + `packages/agent/src`, `packages/coding-agent/src/core` and `src/utils`) in + [Maple #1389](https://github.com/MaplePrivacyLabs/Maple/pull/1389): + 22 upstream commits adopted, among them dropping another model's thinking, + thinking-free summaries without a fixed output cap, and `before_agent_start` + for custom-message runs; provider, catalog, terminal and unported features + skipped. Upstream's session manager and tool output did not change, so the + storage adaptations stay. Skipped `bd3c9b823` (activating extension tools + registered after a branch recorded its tools): it needs every tool the + branch ever named, which a `SessionBase` does not carry. Validated with + `just ci`. diff --git a/apps/maple-agent/crates/maple-agent/src/agent/tools/tests.rs b/apps/maple-agent/crates/maple-agent/src/agent/tools/tests.rs index 0434c3f47..0ef625f28 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent/tools/tests.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent/tools/tests.rs @@ -235,6 +235,7 @@ fn near_miss_questions_take_the_schema_shape() { ask.declaration(), &pi_ai::ToolCall { id: "call-1".to_string(), + arguments_error: None, partial_json: None, name: "request_user_input".to_string(), arguments: match expected { diff --git a/apps/maple-agent/crates/maple-agent/src/store/paging/tests.rs b/apps/maple-agent/crates/maple-agent/src/store/paging/tests.rs index 411ebf4ac..3c76a7999 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/paging/tests.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/paging/tests.rs @@ -640,6 +640,7 @@ fn latest_plan_is_returned_even_when_it_precedes_the_tail_page() { .as_object() .unwrap() .clone(), + arguments_error: None, partial_json: None, })], api: "test".into(), diff --git a/apps/maple-agent/crates/maple-agent/src/store/projector/tests.rs b/apps/maple-agent/crates/maple-agent/src/store/projector/tests.rs index dc4b8d2ec..837ea7ebf 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/projector/tests.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/projector/tests.rs @@ -49,6 +49,7 @@ fn assistant(timestamp: i64, call: Option<&str>) -> EntryKind { id: id.into(), name: "read".into(), arguments: json!({"path":"README.md"}).as_object().unwrap().clone(), + arguments_error: None, partial_json: None, })); } @@ -485,6 +486,7 @@ fn chunked_rebuild_resumes_inside_a_large_entry_and_publishes_once() { id: format!("call-{index}"), name: "read".into(), arguments: json!({"path":"file"}).as_object().unwrap().clone(), + arguments_error: None, partial_json: None, })); } diff --git a/apps/maple-agent/crates/maple-agent/src/store/replay/tests.rs b/apps/maple-agent/crates/maple-agent/src/store/replay/tests.rs index 73e446902..d5b568b1a 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/replay/tests.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/replay/tests.rs @@ -81,6 +81,7 @@ fn call(input: &str) -> SessionMessage { .as_object() .unwrap() .clone(), + arguments_error: None, partial_json: None, })], api: "fixture".into(), diff --git a/apps/maple-agent/crates/pi-agent-core/tests/agent_loop.rs b/apps/maple-agent/crates/pi-agent-core/tests/agent_loop.rs index a6f820f0a..01b0e2d2f 100644 --- a/apps/maple-agent/crates/pi-agent-core/tests/agent_loop.rs +++ b/apps/maple-agent/crates/pi-agent-core/tests/agent_loop.rs @@ -443,6 +443,46 @@ async fn invalid_arguments_unknown_tools_and_panics_become_error_results() { ); } +#[tokio::test] +async fn malformed_argument_text_fails_the_call_without_running_it() { + let faux = FauxProvider::new(); + faux.push_text("done"); + // The first reply streams arguments that are not JSON; the faux provider only sends + // valid ones. + let replies = faux.clone(); + let first = AtomicUsize::new(0); + let stream_fn = + move |model: &pi_ai::Model, context: pi_ai::Context, options: pi_ai::StreamOptions| { + if first.fetch_add(1, Ordering::SeqCst) > 0 { + return pi_ai::StreamFn::stream(&replies, model, context, options); + } + let (sender, stream) = pi_ai::AssistantMessageStream::channel(); + let mut builder = pi_ai::AssistantMessageBuilder::new(sender, model); + builder.start(); + let index = builder.tool_call_start("c1", "explode"); + builder.tool_call_delta(index, "{\"a\":.5}"); + builder.finish(StopReason::ToolUse, pi_ai::Usage::default()); + stream + }; + // The tool panics if it runs. + let explode = Arc::new(Explode { + declaration: Tool::new("explode", "Panics", json!({ "type": "object" })), + }); + let config = AgentLoopConfig::new(faux.model(), Arc::new(stream_fn)); + let (messages, _) = run(&faux, config, context(vec![explode]), "go").await; + let text = result_text(&messages[3]); + assert!( + text.starts_with("Invalid arguments for tool \"explode\": ") + && text.ends_with("\n\nReceived arguments:\n{\"a\":.5}"), + "{text}" + ); + assert!(matches!(&messages[3], Message::ToolResult(result) if result.is_error)); + assert_eq!( + roles(&messages), + ["system", "user", "assistant", "toolResult", "assistant"] + ); +} + fn sleeper( name: &str, millis: u64, @@ -890,6 +930,7 @@ async fn run_tool_call_applies_the_same_gates_outside_the_loop() { id: "nested/1".into(), name: "echo".into(), arguments: json!({ "text": "hi" }).as_object().cloned().unwrap(), + arguments_error: None, partial_json: None, }; let allowed = run_tool_call( diff --git a/apps/maple-agent/crates/pi-ai/src/json.rs b/apps/maple-agent/crates/pi-ai/src/json.rs index b417eb51c..2576124ba 100644 --- a/apps/maple-agent/crates/pi-ai/src/json.rs +++ b/apps/maple-agent/crates/pi-ai/src/json.rs @@ -25,6 +25,31 @@ pub fn parse_streaming_json(raw: &str) -> Map { .unwrap_or_default() } +/// Parse the complete argument text of a finished tool call, as Pi's +/// `parseToolCallArguments`. Unlike [`parse_streaming_json`], text that is not a JSON +/// object, even after its strings are repaired, is not cut back to a valid prefix: the +/// error describes it and includes the text, so validation can show the model what it +/// sent. +pub fn parse_tool_call_arguments(raw: &str) -> Result, String> { + if raw.trim().is_empty() { + return Ok(Map::new()); + } + let parsed = serde_json::from_str::(raw).or_else(|error| { + let repaired = repair_strings(raw); + if repaired == raw { + return Err(error); + } + serde_json::from_str::(&repaired) + }); + match parsed { + Ok(Value::Object(map)) => Ok(map), + Ok(_) => Err(format!( + "Tool arguments must be a JSON object\n\nReceived arguments:\n{raw}" + )), + Err(error) => Err(format!("{error}\n\nReceived arguments:\n{raw}")), + } +} + /// Escape what JSON does not allow inside strings: raw control characters, and /// backslashes that start no valid escape. fn repair_strings(json: &str) -> String { @@ -51,7 +76,8 @@ fn repair_strings(json: &str) -> String { repaired.push_str(&json[offset..offset + 6]); chars.nth(4); } - Some(next @ ('"' | '\\' | '/' | 'b' | 'f' | 'n' | 'r' | 't' | 'u')) => { + // `\u` without four hex digits, as in `\usepackage`, is no escape either. + Some(next @ ('"' | '\\' | '/' | 'b' | 'f' | 'n' | 'r' | 't')) => { repaired.push('\\'); repaired.push(next); chars.next(); @@ -275,6 +301,51 @@ mod tests { ); } + #[test] + fn a_backslash_u_without_hex_digits_stays_literal() { + assert_eq!( + parse(r#"{"text":"\url{x}"}"#), + json!({ "text": "\\url{x}" }) + ); + // The fields after it are kept. + assert_eq!( + parse(r#"{"path":"a.tex","content":"\usepackage{x}","mode":"w"}"#), + json!({ "path": "a.tex", "content": "\\usepackage{x}", "mode": "w" }) + ); + assert_eq!( + parse(r#"{"text":"caf\u00e9 \url"}"#), + json!({ "text": "caf\u{e9} \\url" }) + ); + assert_eq!(parse(r#"{"text":"ab\u00"#), json!({ "text": "ab" })); + } + + #[test] + fn finished_arguments_are_parsed_strictly() { + assert_eq!( + parse_tool_call_arguments(r#"{"kind":"seek","note":"\url"}"#), + Ok(json!({ "kind": "seek", "note": "\\url" }) + .as_object() + .cloned() + .unwrap()) + ); + assert_eq!(parse_tool_call_arguments(" "), Ok(Map::new())); + // Neither the valid prefix of malformed JSON nor cut-off JSON is kept. + for raw in [ + r#"{"action":"apply","actions":[{"kind":"seek","time":.5}]}"#, + r#"{"kind":"seek","clipId":"clip-fl"#, + ] { + let error = parse_tool_call_arguments(raw).unwrap_err(); + assert!( + error.ends_with(&format!("\n\nReceived arguments:\n{raw}")), + "{error}" + ); + } + assert_eq!( + parse_tool_call_arguments("[1,2]"), + Err("Tool arguments must be a JSON object\n\nReceived arguments:\n[1,2]".into()) + ); + } + #[test] fn non_objects_and_garbage_become_empty() { assert_eq!(parse(""), json!({})); diff --git a/apps/maple-agent/crates/pi-ai/src/lib.rs b/apps/maple-agent/crates/pi-ai/src/lib.rs index d1a7c016a..fac6ac6d7 100644 --- a/apps/maple-agent/crates/pi-ai/src/lib.rs +++ b/apps/maple-agent/crates/pi-ai/src/lib.rs @@ -37,7 +37,7 @@ pub use classifier::{ ClassifierStopReason, ClassifyFn, }; pub use estimate::{estimate_message_tokens, estimate_text_tokens, estimate_tool_tokens}; -pub use json::parse_streaming_json; +pub use json::{parse_streaming_json, parse_tool_call_arguments}; pub use overflow::{ is_context_overflow, is_recoverable_length, is_retryable_error, retry_delay_ms, }; diff --git a/apps/maple-agent/crates/pi-ai/src/openai.rs b/apps/maple-agent/crates/pi-ai/src/openai.rs index 8dc563678..5abafc6d0 100644 --- a/apps/maple-agent/crates/pi-ai/src/openai.rs +++ b/apps/maple-agent/crates/pi-ai/src/openai.rs @@ -600,7 +600,8 @@ fn convert_messages(model: &Model, context: &Context) -> Vec { } Message::Assistant(assistant) => { // A model gets its own reasoning back, in the field it came in. Another - // model's reasoning goes back as plain text. + // model's reasoning is dropped: sent as text, it reads as part of the reply, + // and the new model can start writing its reasoning into its answers. let same_model = assistant.provider == model.provider && assistant.api == model.api && assistant.model == model.id; @@ -613,15 +614,13 @@ fn convert_messages(model: &Model, context: &Context) -> Vec { text_parts.push(text.text.as_str()) } AssistantContent::Thinking(thinking) - if !thinking.redacted && !thinking.thinking.trim().is_empty() => + if same_model + && !thinking.redacted + && !thinking.thinking.trim().is_empty() => { - if same_model { - reasoning.push(thinking.thinking.as_str()); - reasoning_field = - reasoning_field.or(thinking.thinking_signature.as_deref()); - } else { - text_parts.push(thinking.thinking.as_str()); - } + reasoning.push(thinking.thinking.as_str()); + reasoning_field = + reasoning_field.or(thinking.thinking_signature.as_deref()); } _ => {} } @@ -664,7 +663,9 @@ fn convert_messages(model: &Model, context: &Context) -> Vec { if has_images { images_from_tools.extend(content.iter().filter_map(image_part)); } - let text = match (text.is_empty(), has_images) { + // Some servers reject output that is only whitespace, such as a bare + // line break, so it counts as none. + let text = match (text.trim().is_empty(), has_images) { (false, _) => text, (true, true) => "(see attached image)".to_string(), (true, false) => "(no tool output)".to_string(), @@ -1165,7 +1166,7 @@ impl HttpTransport for ReqwestTransport { mod tests { use super::*; use crate::overflow::{is_context_overflow, is_retryable_error}; - use crate::types::{SystemMessage, ThinkingContent, Tool, ToolChoice, UserMessage}; + use crate::types::{SystemMessage, ThinkingContent, Tool, ToolCall, ToolChoice, UserMessage}; use futures_util::FutureExt; use std::sync::Mutex; @@ -1296,6 +1297,44 @@ mod tests { ); } + #[tokio::test] + async fn malformed_final_arguments_are_reported_instead_of_a_parsed_prefix() { + let transport = transport( + 200, + vec![ + "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"c1\",\"function\":{\"name\":\"read\",\"arguments\":\"{\\\"path\\\":\\\"a\\\",\\\"offset\\\":.5}\"}}]}}]}\n\n", + "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":1,\"id\":\"c2\",\"function\":{\"name\":\"read\",\"arguments\":\"{\\\"path\\\":\\\"\\\\usepackage\\\"}\"}}]},\"finish_reason\":\"tool_calls\"}]}\n\n", + "data: [DONE]\n\n", + ], + ); + let provider = OpenAiCompletions::new(transport); + let context = Context::new( + "sys", + vec![tool()], + vec![Message::User(UserMessage::text("hi"))], + ); + let options = StreamOptions { + api_key: Some("key".into()), + ..StreamOptions::default() + }; + let message = provider.stream(&model(), context, options).result().await; + let calls: Vec<&ToolCall> = message.tool_calls().collect(); + assert!(calls[0].arguments.is_empty()); + let error = calls[0].arguments_error.as_deref().unwrap(); + assert!( + error.ends_with("\n\nReceived arguments:\n{\"path\":\"a\",\"offset\":.5}"), + "{error}" + ); + // A repairable escape is no error. + assert_eq!(calls[1].arguments["path"], "\\usepackage"); + assert_eq!(calls[1].arguments_error, None); + let error = crate::validate_tool_arguments(&tool(), calls[0]).unwrap_err(); + assert!( + error.starts_with("Invalid arguments for tool \"read\": "), + "{error}" + ); + } + #[test] fn tools_that_ask_for_it_go_in_strict_mode_where_the_provider_has_it() { use crate::types::{ConstrainedSampling, StrictSampling}; @@ -2290,22 +2329,93 @@ mod tests { } #[test] - fn another_models_reasoning_goes_back_as_text() { + fn another_models_reasoning_is_dropped() { let model = model(); let mut earlier = AssistantMessage::empty(&model); earlier.model = "other-model".into(); earlier.content = vec![ - thinking("plan", Some("reasoning_content")), - AssistantContent::text("answer"), + thinking("Candidates: proxy, IPv6, TLS.", Some("reasoning_content")), + AssistantContent::text("Let me check the TLS versions."), + AssistantContent::tool_call("c1", "bash", json!({ "command": "curl -v" })), ]; + let mut thought_only = AssistantMessage::empty(&model); + thought_only.model = "other-model".into(); + thought_only.content = vec![thinking("nothing to say", Some("reasoning"))]; let context = Context::from_messages(vec![ Message::User(UserMessage::text("a")), Message::Assistant(earlier), + Message::ToolResult(ToolResultMessage { + tool_call_id: "c1".into(), + tool_name: "bash".into(), + content: vec![Content::text("TLSv1.2")], + details: None, + usage: None, + is_error: false, + timestamp: 0, + }), + Message::User(UserMessage::text("b")), + Message::Assistant(thought_only), + Message::User(UserMessage::text("c")), ]); let messages = convert_messages(&model, &context); - // Text blocks join without a separator. - assert_eq!(messages[1]["content"], "plananswer"); - assert!(messages[1].get("reasoning_content").is_none()); + // The text and the tool call stay; the reasoning goes in no field and not as text. + assert_eq!(messages[1]["content"], "Let me check the TLS versions."); + assert_eq!(messages[1]["tool_calls"][0]["id"], "c1"); + assert!( + REASONING_FIELDS + .iter() + .all(|field| messages[1].get(*field).is_none()) + ); + // A reply that was only reasoning leaves nothing to send. + let roles: Vec<&str> = messages + .iter() + .map(|message| message["role"].as_str().unwrap()) + .collect(); + assert_eq!(roles, ["user", "assistant", "tool", "user", "user"]); + } + + #[test] + fn only_the_target_models_reasoning_comes_back_when_models_alternate() { + let model = model(); + let reply = |id: &str, thought: &str, text: &str| { + let mut reply = AssistantMessage::empty(&model); + reply.model = id.into(); + reply.content = vec![ + thinking(thought, Some("reasoning_content")), + AssistantContent::text(text), + ]; + Message::Assistant(reply) + }; + let context = Context::from_messages(vec![ + Message::User(UserMessage::text("1")), + reply(&model.id, "mine 1", "a"), + Message::User(UserMessage::text("2")), + reply("other-model", "theirs 1", "b"), + Message::User(UserMessage::text("3")), + reply(&model.id, "mine 2", "c"), + Message::User(UserMessage::text("4")), + reply("other-model", "theirs 2", "d"), + Message::User(UserMessage::text("hi")), + ]); + let replies: Vec<(Value, Option)> = convert_messages(&model, &context) + .into_iter() + .filter(|message| message["role"] == "assistant") + .map(|message| { + ( + message["content"].clone(), + message.get("reasoning_content").cloned(), + ) + }) + .collect(); + assert_eq!( + replies, + [ + (json!("a"), Some(json!("mine 1"))), + (json!("b"), None), + (json!("c"), Some(json!("mine 2"))), + (json!("d"), None), + ] + ); } #[test] @@ -2405,6 +2515,20 @@ mod tests { assert_eq!(messages[4]["content"], "(no tool output)"); } + #[test] + fn whitespace_only_tool_output_counts_as_none() { + let model = model(); + let messages = tool_exchange(vec![ + vec![Content::text("\r\n")], + vec![Content::text(" \t"), Content::image("AAAA", "image/png")], + vec![Content::text(" kept as is \n")], + ]); + let messages = convert_messages(&model, &Context::from_messages(messages)); + assert_eq!(messages[2]["content"], "(no tool output)"); + assert_eq!(messages[3]["content"], "(see attached image)"); + assert_eq!(messages[4]["content"], " kept as is \n"); + } + #[test] fn utf8_split_across_chunks_is_reassembled() { let mut decoder = Utf8Decoder::default(); diff --git a/apps/maple-agent/crates/pi-ai/src/overflow.rs b/apps/maple-agent/crates/pi-ai/src/overflow.rs index fe3e6134d..657038e30 100644 --- a/apps/maple-agent/crates/pi-ai/src/overflow.rs +++ b/apps/maple-agent/crates/pi-ai/src/overflow.rs @@ -32,6 +32,9 @@ static OVERFLOW: LazyLock = LazyLock::new(|| { r"(?i)context[_ ]length[_ ]exceeded", r"(?i)too many tokens", r"(?i)token limit exceeded", + // An HTTP 413 is a request too large for the provider, whatever its body says. + // Errors start with the status, or with a provider's own prefix. + r"(?i)^(?:[\w -]+ API error \(413\):|413(?:\s|:|$))", ]) .expect("overflow patterns compile") }); @@ -96,6 +99,7 @@ static ACCOUNT_LIMIT: LazyLock = LazyLock::new(|| { r"(?i)GoUsageLimitError|FreeUsageLimitError|Monthly usage limit reached|available balance", r"|insufficient_quota|out of budget|quota exceeded|billing", r"|subscription_sharing_usage_limit_exceeded", + r"|subscription_sharing_workplace_individual_usage_limit_exceeded", )) .expect("pattern compiles") }); @@ -104,13 +108,20 @@ static RETRYABLE: LazyLock = LazyLock::new(|| { Regex::new(concat!( r"(?i)overloaded|server_busy|servers are currently busy|high demand|at capacity", r"|rate.?limit|too many requests", - r"|\b(?:429|500|502|503|504|520|524)\b|service.?unavailable|server.?error|internal.?error", - r"|provider.?returned.?error|exceeded request buffer limit while retrying upstream", + // A status code, but not digits inside a longer word or number, such as a + // request id or a byte count. + r"|(?:^|[^a-z0-9])(?:429|500|502|503|504|520|524)(?:[^a-z0-9]|$)", + r"|service.?(?:is.?)?(?:temporarily.?)?unavailable|server.?error|internal.?error", + r"|provider.?returned.?(?:an.?)?error|exceeded request buffer limit while retrying upstream", + r"|unable to forward upstream|upstream.?request.?failed|upstream_stream_read_error", + r"|stream was interrupted|http/?2 stream failed", r"|network.?error|connection.?(?:error|refused|lost|reset)", - r"|other side closed|fetch failed|getaddrinfo|ENOTFOUND|EAI_AGAIN|upstream.?connect", + r"|other side closed|fetch failed|request to \S+ failed, reason:", + r"|getaddrinfo|ENOTFOUND|EAI_AGAIN|upstream.?connect", r"|reset before headers|socket hang up|socket connection was closed|timed? ?out", r"|terminated|websocket.?(?:closed|error)", r"|ended without|stream ended|http2 request did not get a response", + r"|stream disconnected before completion", r"|pending stream has been canceled|retry delay", r"|retry your request|try your request again|ResourceExhausted", r"|subscription_sharing_(?:usage|user)_unavailable", @@ -207,14 +218,15 @@ mod tests { } #[test] - fn bodyless_400_and_413_errors_are_overflow_only_from_cerebras() { - for error in ["400 status code (no body)", "413 status code (no body)"] { - let mut message = failed(error); - message.provider = "cerebras".into(); - assert!(is_context_overflow(&message, Some(131_072)), "{error}"); - message.provider = "opencode-go".into(); - assert!(!is_context_overflow(&message, Some(1_000_000)), "{error}"); - } + fn a_bodyless_400_is_overflow_only_from_cerebras() { + let mut message = failed("400 status code (no body)"); + message.provider = "cerebras".into(); + assert!(is_context_overflow(&message, Some(131_072))); + message.provider = "opencode-go".into(); + assert!(!is_context_overflow(&message, Some(1_000_000))); + // A 413 is overflow from every provider. + message.error_message = Some("413 status code (no body)".into()); + assert!(is_context_overflow(&message, Some(1_000_000))); } #[test] @@ -279,14 +291,55 @@ mod tests { "Request timed out", "{\"code\":\"server_busy\",\"message\":\"Busy\"}", "Our servers are currently busy, please come back later", + "The service is temporarily unavailable", + "The upstream provider returned an error", + "Unable to forward upstream stream", + "upstream_error: Upstream request failed", + "upstream_stream_read_error", + "The stream was interrupted", + "Upstream HTTP/2 stream failed", + "request to https://oauth2.googleapis.com/token failed, reason: ", + "stream disconnected before completion: stream closed before response.completed", + "status_503", ] { assert!(is_retryable_error(&failed(error)), "{error}"); } - for error in ["401 invalid api key", "400 bad request: unknown field"] { + for error in [ + "401 invalid api key", + "400 bad request: unknown field", + // Digits inside a request id or a byte count are no status code. + "400 invalid request req_011CU503xyz", + "400 the body exceeds 5242880 bytes", + ] { assert!(!is_retryable_error(&failed(error)), "{error}"); } } + #[test] + fn a_413_is_overflow_whatever_its_body_says() { + for error in [ + "413 {\"message\":\"failed to parse request\",\"code\":\"\"}", + "413", + "413: Payload Too Large", + "github-copilot API error (413): {\"message\":\"failed to parse request\"}", + ] { + assert!(is_context_overflow(&failed(error), None), "{error}"); + } + // Only the status, not a number in the body or a longer one. + for error in [ + "github-copilot API error (400): {\"message\":\"failed to parse request: invalid value 413\"}", + "custom-provider API error (400): upstream mentioned OpenAI API error (413): invalid input", + "custom-provider API error (4130): invalid input", + "4130 failed to parse request", + ] { + assert!(!is_context_overflow(&failed(error), None), "{error}"); + } + // Only for a failed response. + let mut stopped = failed("413 Payload Too Large"); + stopped.stop_reason = StopReason::Aborted; + assert!(!is_context_overflow(&stopped, None)); + } + #[test] fn exhausted_quotas_are_not_retried_even_as_429s() { for error in [ @@ -295,6 +348,7 @@ mod tests { "Rate limit: quota exceeded for this key", "402 billing hard limit reached", "subscription_sharing_usage_limit_exceeded", + "429 subscription_sharing_workplace_individual_usage_limit_exceeded", ] { assert!(!is_retryable_error(&failed(error)), "{error}"); } diff --git a/apps/maple-agent/crates/pi-ai/src/stream.rs b/apps/maple-agent/crates/pi-ai/src/stream.rs index a87cc85b9..a156ef2b6 100644 --- a/apps/maple-agent/crates/pi-ai/src/stream.rs +++ b/apps/maple-agent/crates/pi-ai/src/stream.rs @@ -6,7 +6,7 @@ use serde::{Deserialize, Serialize}; use serde_json::Map; use tokio::sync::mpsc; -use crate::json::parse_streaming_json; +use crate::json::{parse_streaming_json, parse_tool_call_arguments}; use crate::types::{ AssistantContent, AssistantMessage, Model, StopReason, TextContent, ThinkingContent, ToolCall, Usage, now_ms, @@ -90,6 +90,7 @@ pub fn apply_event(partial: &mut AssistantMessage, event: &AssistantMessageEvent id: id.clone(), name: name.clone(), arguments: Map::new(), + arguments_error: None, partial_json: Some(String::new()), })) } @@ -269,7 +270,7 @@ impl AssistantMessageBuilder { let index = match self.open { Some(OpenBlock::Text(index)) => index, _ => { - self.close_block(); + self.close_block(false); self.start(); let index = self.partial.content.len(); self.partial.content.push(AssistantContent::text("")); @@ -296,7 +297,7 @@ impl AssistantMessageBuilder { let index = match self.open { Some(OpenBlock::Thinking(index)) => index, _ => { - self.close_block(); + self.close_block(false); self.start(); let index = self.partial.content.len(); self.partial.content.push(AssistantContent::thinking("")); @@ -318,7 +319,7 @@ impl AssistantMessageBuilder { /// Open a tool-call block and return its content index. pub fn tool_call_start(&mut self, id: impl Into, name: impl Into) -> usize { - self.close_block(); + self.close_block(false); self.start(); let index = self.partial.content.len(); let (id, name) = (id.into(), name.into()); @@ -328,6 +329,7 @@ impl AssistantMessageBuilder { id: id.clone(), name: name.clone(), arguments: Map::new(), + arguments_error: None, partial_json: None, })); self.arguments.push((index, String::new())); @@ -363,7 +365,12 @@ impl AssistantMessageBuilder { } } - fn finish_tool_call(&mut self, index: usize) -> Option { + /// Parse the arguments of the tool call at `index`. A call closes when the next block + /// starts and can still get fragments, so until the response finishes its arguments + /// are read as they stream. Once it finishes, as when Pi's stream ends, the complete + /// text is parsed strictly: malformed arguments leave no repaired prefix, and the call + /// carries the error instead, which validation reports to the model. + fn finish_tool_call(&mut self, index: usize, complete: bool) -> Option { let raw = self .arguments .iter() @@ -372,18 +379,25 @@ impl AssistantMessageBuilder { let AssistantContent::ToolCall(call) = self.partial.content.get_mut(index)? else { return None; }; - call.arguments = parse_streaming_json(&raw); + (call.arguments, call.arguments_error) = if complete { + match parse_tool_call_arguments(&raw) { + Ok(arguments) => (arguments, None), + Err(error) => (Map::new(), Some(error)), + } + } else { + (parse_streaming_json(&raw), None) + }; Some(call.clone()) } - fn close_block(&mut self) { + fn close_block(&mut self, complete: bool) { match self.open.take() { Some(OpenBlock::Text(index)) => self.send(AssistantMessageEvent::TextEnd { index }), Some(OpenBlock::Thinking(index)) => { self.send(AssistantMessageEvent::ThinkingEnd { index }) } Some(OpenBlock::ToolCall(index)) => { - if let Some(tool_call) = self.finish_tool_call(index) { + if let Some(tool_call) = self.finish_tool_call(index, complete) { self.send(AssistantMessageEvent::ToolCallEnd { index, tool_call }); } } @@ -393,10 +407,10 @@ impl AssistantMessageBuilder { /// Re-parse every tool call's arguments, including fragments that arrived for a block /// after it was closed. - fn settle_tool_calls(&mut self) { + fn settle_tool_calls(&mut self, complete: bool) { let indexes: Vec = self.arguments.iter().map(|(index, _)| *index).collect(); for index in indexes { - self.finish_tool_call(index); + self.finish_tool_call(index, complete); } } @@ -415,9 +429,9 @@ impl AssistantMessageBuilder { /// End the response successfully and return the final message. pub fn finish(mut self, stop_reason: StopReason, usage: Usage) -> AssistantMessage { - self.close_block(); + self.close_block(true); self.start(); - self.settle_tool_calls(); + self.settle_tool_calls(true); self.partial.stop_reason = stop_reason; self.partial.usage = usage; let message = self.partial.clone(); @@ -429,8 +443,8 @@ impl AssistantMessageBuilder { /// End the response with an error or abort, keeping the content received so far. pub fn fail(mut self, stop_reason: StopReason, error: impl Into) -> AssistantMessage { - self.close_block(); - self.settle_tool_calls(); + self.close_block(false); + self.settle_tool_calls(false); self.partial.stop_reason = stop_reason; self.partial.error_message = Some(error.into()); let message = self.partial.clone(); @@ -513,6 +527,32 @@ mod tests { assert_eq!(partial.content, message.content); } + #[tokio::test] + async fn arguments_cut_off_when_the_response_ends_carry_an_error() { + let (sender, stream) = AssistantMessageStream::channel(); + let mut builder = AssistantMessageBuilder::new(sender, &model()); + let call = builder.tool_call_start("c1", "write"); + builder.tool_call_delta(call, "{\"path\":\"a.txt\",\"content\":\"hel"); + let message = builder.finish(StopReason::Length, Usage::default()); + let call = message.tool_calls().next().unwrap(); + assert!(call.arguments.is_empty()); + assert!( + call.arguments_error + .as_deref() + .unwrap() + .ends_with("\n\nReceived arguments:\n{\"path\":\"a.txt\",\"content\":\"hel") + ); + // The call's end event says the same. + let ended = collect(stream) + .await + .into_iter() + .find_map(|event| match event { + AssistantMessageEvent::ToolCallEnd { tool_call, .. } => Some(tool_call), + _ => None, + }); + assert_eq!(ended.as_ref(), Some(call)); + } + #[tokio::test] async fn applied_events_parse_a_calls_arguments_as_they_stream() { let (sender, stream) = AssistantMessageStream::channel(); diff --git a/apps/maple-agent/crates/pi-ai/src/types.rs b/apps/maple-agent/crates/pi-ai/src/types.rs index 61205548f..dc5980911 100644 --- a/apps/maple-agent/crates/pi-ai/src/types.rs +++ b/apps/maple-agent/crates/pi-ai/src/types.rs @@ -67,6 +67,10 @@ pub struct ToolCall { pub name: String, #[serde(default)] pub arguments: Map, + /// Why the call's complete argument text was not a JSON object, with that text, as + /// Pi's `argumentsError`. `arguments` is then empty, and validation fails with this. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub arguments_error: Option, /// The argument JSON received so far, while [`crate::apply_event`] follows the call /// as it streams; `arguments` holds its parse. `None` once the call has ended. Never /// saved. @@ -139,6 +143,7 @@ impl AssistantContent { id: id.into(), name: name.into(), arguments, + arguments_error: None, partial_json: None, }) } diff --git a/apps/maple-agent/crates/pi-ai/src/validation.rs b/apps/maple-agent/crates/pi-ai/src/validation.rs index 9b982069b..19de988d1 100644 --- a/apps/maple-agent/crates/pi-ai/src/validation.rs +++ b/apps/maple-agent/crates/pi-ai/src/validation.rs @@ -10,6 +10,12 @@ use crate::types::{Tool, ToolCall}; /// `anyOf` and `oneOf` too. Returns the arguments to execute with, or an error message /// written for the model. pub fn validate_tool_arguments(tool: &Tool, call: &ToolCall) -> Result { + if let Some(error) = &call.arguments_error { + return Err(format!( + "Invalid arguments for tool \"{}\": {error}", + call.name + )); + } let mut args = Value::Object(call.arguments.clone()); normalize_optional_nulls(&mut args, &tool.parameters); coerce(&mut args, &tool.parameters); @@ -290,6 +296,7 @@ mod tests { id: "c1".into(), name: "read".into(), arguments: arguments.as_object().cloned().unwrap(), + arguments_error: None, partial_json: None, } } @@ -300,6 +307,16 @@ mod tests { assert_eq!(args, json!({ "path": "a" })); } + #[test] + fn malformed_argument_text_fails_with_its_error() { + let mut malformed = call(json!({})); + malformed.arguments_error = Some("Unexpected token".into()); + assert_eq!( + validate_tool_arguments(&tool(), &malformed), + Err("Invalid arguments for tool \"read\": Unexpected token".into()) + ); + } + #[test] fn primitives_are_coerced_to_the_schema_type() { let args = validate_tool_arguments( diff --git a/apps/maple-agent/crates/pi-coding-agent/src/agent_session.rs b/apps/maple-agent/crates/pi-coding-agent/src/agent_session.rs index 1d7533d5f..482f44123 100644 --- a/apps/maple-agent/crates/pi-coding-agent/src/agent_session.rs +++ b/apps/maple-agent/crates/pi-coding-agent/src/agent_session.rs @@ -45,11 +45,11 @@ use crate::extensions::{ CustomMessageDraft, Extension, ExtensionContext, ExtensionErrorReport, ExtensionRunner, ExtensionUi, ForkPosition, Input, InputAction, InputSource, MessageEnd, ModelSelect, ProviderStreamEvent, RegisteredTool, ResourcesDiscover, ResourcesDiscoverReason, - ResourcesDiscoverResult, SessionBeforeCompact, SessionBeforeFork, SessionBeforeSwitch, - SessionBeforeTree, SessionCompact, SessionCompactFailed, SessionInfoChanged, SessionShutdown, - SessionStart, SessionStartReason, SessionTree, SlashCommandInfo, SlashCommandSource, - ThinkingLevelSelect, ToolCall, ToolInfo, ToolPrompt, ToolResult, TreePreparation, TreeSummary, - TurnEnd, UserBash, UserBashResult, + ResourcesDiscoverResult, RunTrigger, SessionBeforeCompact, SessionBeforeFork, + SessionBeforeSwitch, SessionBeforeTree, SessionCompact, SessionCompactFailed, + SessionInfoChanged, SessionShutdown, SessionStart, SessionStartReason, SessionTree, + SlashCommandInfo, SlashCommandSource, ThinkingLevelSelect, ToolCall, ToolInfo, ToolPrompt, + ToolResult, TreePreparation, TreeSummary, TurnEnd, UserBash, UserBashResult, }; use crate::messages::{BashExecutionMessage, CustomMessage, SessionMessage, convert_to_llm}; use crate::models::ModelRegistry; @@ -1047,8 +1047,9 @@ impl AgentSession { /// Start a turn with a custom message and wait for it, as Pi's `sendCustomMessage` /// does with `triggerTurn`, for a host that hands the model something the user did /// not type, such as a background agent's result. A message not to `display` stays - /// out of the transcript the user sees but reaches the model. During a prompt it is - /// queued as steering instead. + /// out of the transcript the user sees but reaches the model. The run starts as a + /// prompt's does, with `before_agent_start` and the next-turn messages. During a + /// prompt it is queued as steering instead. pub async fn send_custom_message( &self, draft: CustomMessageDraft, @@ -1068,7 +1069,8 @@ impl AgentSession { _ => return Err(AgentSessionError::Busy), } }; - core.run_delivered(activity, custom_message(draft)).await?; + core.start_run(activity, RunTrigger::Custom(custom(draft))) + .await?; Ok(PromptOutcome::Completed) } @@ -1720,14 +1722,18 @@ fn user_message(text: &str, images: Vec) -> SessionMessage { })) } -fn custom_message(draft: CustomMessageDraft) -> SessionMessage { - SessionMessage::Custom(CustomMessage { +fn custom(draft: CustomMessageDraft) -> CustomMessage { + CustomMessage { custom_type: draft.custom_type, content: draft.content, display: draft.display, details: draft.details, timestamp: now_ms(), - }) + } +} + +fn custom_message(draft: CustomMessageDraft) -> SessionMessage { + SessionMessage::Custom(custom(draft)) } impl SessionCore { @@ -2288,11 +2294,47 @@ impl SessionCore { _ => return Err(AgentSessionError::Busy), } }; + self.start_run( + activity, + RunTrigger::User { + prompt: text, + images, + }, + ) + .await?; + Ok(PromptOutcome::Completed) + } + + /// Start a run from a prompt or a custom message, as Pi's `_startRun`. The session is + /// busy from here on, so messages sent while `before_agent_start` handlers run join + /// this run, and an abort requested meanwhile stops it once it starts. + async fn start_run( + self: &Arc, + activity: ActivityGuard<'_>, + trigger: RunTrigger, + ) -> Result<(), AgentSessionError> { + lock(&self.run).abort_requested = false; + let prepared = self.prepare_run(trigger).await; + // What extensions and the user's own commands added meanwhile goes before the + // run's messages, as when idle; no run exists yet to split a call from its result. + self.flush_pending_bash(); + self.flush_pending_custom(); + self.run_prompt(activity, prepared?).await + } + + /// The messages that start a run, as Pi's `_prepareRun`: compact first when needed, + /// ask `before_agent_start`, apply the run's prompt options and tools, then the + /// starting message, the waiting next-turn messages and the handlers' messages, after + /// a prompt update when the prompt changed. + async fn prepare_run( + self: &Arc, + trigger: RunTrigger, + ) -> Result, AgentSessionError> { if self.model().is_none() { return Err(AgentSessionError::NoModel); } // As in Pi, a stopped run, or a switch to a smaller model, can leave a context to - // compact before the new prompt, which is the one that runs next. + // compact before the new message, which is the one that runs next. let last = self.agent.with_messages(|messages| { messages.iter().rev().find_map(|message| match message { SessionMessage::Llm(Message::Assistant(assistant)) => Some(assistant.clone()), @@ -2305,8 +2347,7 @@ impl SessionCore { let base = self.prompt_options(); let event = BeforeAgentStart { - prompt: text.clone(), - images: images.clone(), + trigger: trigger.clone(), system_prompt: render(&base).unwrap_or_default(), options: base.clone(), }; @@ -2325,22 +2366,28 @@ impl SessionCore { .system_update(&self.run_prompt_options())? .into_iter() .collect(); - // As in Pi, attached images are made to fit, after the handlers, with notes on - // what was done; one that cannot be used is left out, its note kept. - let auto_resize = lock(&self.settings).images.auto_resize; - let (images, hints) = normalize_prompt_images(images, auto_resize).await; - if !hints.is_empty() { - text = format!("{text}\n\n{}", hints.join("\n")); - } - messages.push(user_message(&text, images)); + match trigger { + RunTrigger::User { prompt, images } => { + // As in Pi, attached images are made to fit, after the handlers, with notes + // on what was done; one that cannot be used is left out, its note kept. + let auto_resize = lock(&self.settings).images.auto_resize; + let (images, hints) = normalize_prompt_images(images, auto_resize).await; + let text = if hints.is_empty() { + prompt + } else { + format!("{prompt}\n\n{}", hints.join("\n")) + }; + messages.push(user_message(&text, images)); + } + RunTrigger::Custom(message) => messages.push(SessionMessage::Custom(message)), + } { let mut run = lock(&self.run); run.forced_prompt = forced.or(prompt_options.force_prompt.clone()); messages.append(&mut run.next_turn_messages); } messages.extend(drafts.into_iter().map(custom_message)); - self.run_prompt(activity, messages).await?; - Ok(PromptOutcome::Completed) + Ok(messages) } /// Expand a skill command, then a prompt template. @@ -2390,23 +2437,6 @@ impl SessionCore { Ok(PromptOutcome::Queued) } - /// Start a run with a message an extension sent, bringing the prompt up to date first. - async fn run_delivered( - self: &Arc, - activity: ActivityGuard<'_>, - message: SessionMessage, - ) -> Result<(), AgentSessionError> { - if self.model().is_none() { - return Err(AgentSessionError::NoModel); - } - let mut messages: Vec = self - .system_update(&self.prompt_options())? - .into_iter() - .collect(); - messages.push(message); - self.run_prompt(activity, messages).await - } - async fn run_prompt( self: &Arc, activity: ActivityGuard<'_>, @@ -2414,7 +2444,6 @@ impl SessionCore { ) -> Result<(), AgentSessionError> { { let mut run = lock(&self.run); - run.abort_requested = false; run.overflow_recovery_attempted = false; run.boundary_changed = false; } @@ -2861,28 +2890,57 @@ impl SessionCore { completed && !self.abort_requested() } - /// Omit the latest assistant message from the context with a context edit. + /// Omit the latest assistant message from the context with a context edit, and the + /// tool results that answer its calls with it, as Pi's `_omitRecoveryAttempt` does: a + /// result without its call would be sent on its own. fn omit_last_assistant(&self) { - let edit = { + let edits = { let mut session = lock(&self.session); - let target = session - .branch() + let targets: Vec = { + let branch = session.branch(); + let Some((position, assistant)) = + branch + .iter() + .enumerate() + .rev() + .find_map(|(position, entry)| match &entry.kind { + EntryKind::Message { + message: SessionMessage::Llm(Message::Assistant(assistant)), + } => Some((position, assistant)), + _ => None, + }) + else { + return; + }; + let calls: Vec<&str> = assistant + .tool_calls() + .map(|call| call.id.as_str()) + .collect(); + let results = branch[position + 1..] + .iter() + .filter_map(|entry| match &entry.kind { + EntryKind::Message { + message: SessionMessage::Llm(Message::ToolResult(result)), + } if calls.contains(&result.tool_call_id.as_str()) => { + Some(entry.id.clone()) + } + _ => None, + }); + std::iter::once(branch[position].id.clone()) + .chain(results) + .collect() + }; + targets .iter() - .rev() - .find_map(|entry| match &entry.kind { - EntryKind::Message { - message: SessionMessage::Llm(Message::Assistant(_)), - } => Some(entry.id.clone()), - _ => None, - }); - target.and_then(|target| { - let id = session.append_context_edit(&target, None).ok()?; - session.entry(&id).cloned() - }) + .filter_map(|target| { + let id = session.append_context_edit(target, None).ok()?; + session.entry(&id).cloned() + }) + .collect::>() }; self.report_persistence(); self.rebuild_messages(); - if let Some(entry) = edit { + for entry in edits { self.emit(AgentSessionEvent::EntryAppended { entry }); } } @@ -3598,7 +3656,7 @@ struct SessionListener { #[async_trait] impl AgentListener for SessionListener { - async fn on_event(&self, event: &AgentEvent, _cancel: &CancellationToken) { + async fn on_event(&self, event: &AgentEvent, cancel: &CancellationToken) { let Some(core) = self.core.upgrade() else { return; }; @@ -3634,7 +3692,15 @@ impl AgentListener for SessionListener { core.emit_queue(); } } - AgentEvent::AgentStart => lock(&core.run).turn_index = 0, + AgentEvent::AgentStart => { + let mut run = lock(&core.run); + run.turn_index = 0; + // An abort requested while the run was starting, during + // `before_agent_start` say, stops it here, before its first request. + if run.abort_requested { + cancel.cancel(); + } + } AgentEvent::TurnEnd { .. } => lock(&core.run).turn_index += 1, AgentEvent::AgentEnd { messages } => { let will_retry = core.will_retry(messages); @@ -3883,7 +3949,8 @@ impl ExtensionContext { /// Add a message to the conversation as `delivery` says. pub fn send_message(&self, draft: CustomMessageDraft, delivery: Delivery) { let Some(core) = self.core() else { return }; - let message = custom_message(draft); + let custom = custom(draft); + let message = SessionMessage::Custom(custom.clone()); let prompting = core.activity() == Activity::Prompt; let behavior = match delivery { Delivery::NextTurn => { @@ -3912,7 +3979,7 @@ impl ExtensionContext { core.queue(behavior, message); return; }; - if let Err(error) = core.run_delivered(activity, message).await { + if let Err(error) = core.start_run(activity, RunTrigger::Custom(custom)).await { core.runner.report(&extension, "send_message", error); } }); diff --git a/apps/maple-agent/crates/pi-coding-agent/src/compaction.rs b/apps/maple-agent/crates/pi-coding-agent/src/compaction.rs index 2d241a106..6a5c02352 100644 --- a/apps/maple-agent/crates/pi-coding-agent/src/compaction.rs +++ b/apps/maple-agent/crates/pi-coding-agent/src/compaction.rs @@ -25,7 +25,8 @@ use crate::session::{ #[serde(rename_all = "camelCase", default)] pub struct CompactionSettings { pub enabled: bool, - /// Room kept free for the response and the summary request. + /// Room kept free in the context window: compaction starts when less is left. Summary + /// requests use the model's output limit, as in Pi. pub reserve_tokens: u64, pub keep_recent_tokens: u64, /// Token settings for particular models, by `provider/model-id`, as Pi's @@ -519,6 +520,11 @@ const BRANCH_SUMMARY_PREAMBLE: &str = "The user explored a different conversatio const TOOL_RESULT_MAX_CHARS: usize = 2_000; /// Render messages as labelled text, so the summarizer reads them instead of continuing them. +/// +/// Thinking is left out, as in Pi. Replayed thinking is often not part of the model's +/// live context (another model's is dropped, for one), and with it the request could be +/// larger than the context it summarizes. Text, tool calls and tool results record what +/// happened. pub fn serialize_conversation(messages: &[Message]) -> String { let mut parts = Vec::new(); for message in messages { @@ -530,25 +536,17 @@ pub fn serialize_conversation(messages: &[Message]) -> String { } } Message::Assistant(assistant) => { - let mut thinking = Vec::new(); - let mut calls = Vec::new(); - for block in &assistant.content { - match block { - AssistantContent::Thinking(block) => thinking.push(block.thinking.as_str()), - AssistantContent::ToolCall(call) => { - let args: Vec = call - .arguments - .iter() - .map(|(key, value)| format!("{key}={value}")) - .collect(); - calls.push(format!("{}({})", call.name, args.join(", "))); - } - AssistantContent::Text(_) => {} - } - } - if !thinking.is_empty() { - parts.push(format!("[Assistant thinking]: {}", thinking.join("\n"))); - } + let calls: Vec = assistant + .tool_calls() + .map(|call| { + let args: Vec = call + .arguments + .iter() + .map(|(key, value)| format!("{key}={value}")) + .collect(); + format!("{}({})", call.name, args.join(", ")) + }) + .collect(); let text = assistant.text(); if !text.is_empty() { parts.push(format!("[Assistant]: {text}")); @@ -585,10 +583,13 @@ fn llm_messages(messages: &[SessionMessage]) -> Vec { /// Run one summary request. Transient failures are retried; a cut-off or tool-calling /// response is an error, since a partial summary must not become a checkpoint. +/// +/// As in Pi, the request sets no output limit: the provider fits the model's own to the +/// request, which is usually much smaller than the context it summarizes, so reasoning +/// does not use up a fixed budget before the summary is written. async fn summarize( summarizer: &Summarizer<'_>, prompt: String, - max_tokens: u64, label: &str, ) -> Result<(String, Usage), String> { let context = Context::new( @@ -596,11 +597,6 @@ async fn summarize( Vec::new(), vec![Message::User(UserMessage::text(prompt))], ); - let max_tokens = if summarizer.model.max_tokens > 0 { - max_tokens.min(summarizer.model.max_tokens) - } else { - max_tokens - }; let report = |event: SummaryRetry| { if let Some(on_retry) = summarizer.on_retry { on_retry(event); @@ -609,7 +605,7 @@ async fn summarize( let mut attempt = 0; loop { let options = StreamOptions { - max_tokens: Some(max_tokens), + max_tokens: None, reasoning: summarizer .options .reasoning @@ -673,7 +669,6 @@ pub async fn compact( summarizer: &Summarizer<'_>, custom_instructions: Option<&str>, ) -> Result { - let reserve = preparation.settings.reserve_tokens; let focus = custom_instructions .map(|instructions| format!("\n\nAdditional focus: {instructions}")) .unwrap_or_default(); @@ -691,7 +686,7 @@ pub async fn compact( None => SUMMARY_PROMPT, }; prompt.push_str(&format!("{instructions}{focus}")); - summarize(summarizer, prompt, reserve * 4 / 5, "Summarization").await + summarize(summarizer, prompt, "Summarization").await }; let (mut summary, usage) = if preparation.turn_prefix_messages.is_empty() { history.await? @@ -706,10 +701,12 @@ pub async fn compact( history.await? }; let conversation = serialize_conversation(&llm_messages(&preparation.turn_prefix_messages)); - let prompt = - format!("# Conversation\n{conversation}\n\n# Instructions\n{TURN_PREFIX_PROMPT}"); + // The focus applies to the turn's start too, as in Pi. + let prompt = format!( + "# Conversation\n{conversation}\n\n# Instructions\n{TURN_PREFIX_PROMPT}{focus}" + ); let (prefix, prefix_usage) = - summarize(summarizer, prompt, reserve / 2, "Turn prefix summarization").await?; + summarize(summarizer, prompt, "Turn prefix summarization").await?; ( format!("{history_text}\n\n---\n\n**Turn Context (split turn):**\n\n{prefix}"), history_usage.add(&prefix_usage), @@ -807,7 +804,7 @@ pub async fn summarize_branch( None => BRANCH_PROMPT.to_string(), }; let prompt = format!("\n{conversation}\n\n\n{instructions}"); - let (text, usage) = summarize(summarizer, prompt, 4_096, "Branch summarization").await?; + let (text, usage) = summarize(summarizer, prompt, "Branch summarization").await?; let (read_files, modified_files) = file_operations.lists(); let files = format_file_operations(&read_files, &modified_files); Ok(BranchSummaryResult { @@ -1234,6 +1231,87 @@ mod tests { } } + #[test] + fn the_conversation_for_a_summary_leaves_thinking_out() { + let reply = Message::Assistant(AssistantMessage { + content: vec![ + AssistantContent::thinking("the user probably means the parser"), + AssistantContent::text("Reading the parser."), + faux_tool_call("read", json!({ "path": "src/parser.rs" })), + ], + ..AssistantMessage::empty(&FauxProvider::default_model()) + }); + let thought_only = Message::Assistant(AssistantMessage { + content: vec![AssistantContent::thinking("nothing to add")], + ..AssistantMessage::empty(&FauxProvider::default_model()) + }); + let text = serialize_conversation(&[ + Message::User(UserMessage::text("fix it")), + reply, + thought_only, + ]); + assert_eq!( + text, + "[User]: fix it\n\n[Assistant]: Reading the parser.\n\n[Assistant tool calls]: read(path=\"src/parser.rs\")" + ); + } + + #[tokio::test] + async fn summary_requests_set_no_output_cap_and_the_turn_start_keeps_the_focus() { + // With history before the split turn, and with only the split turn since the last + // compaction. + for history in [vec![user("an older question")], Vec::new()] { + let preparation = CompactionPreparation { + first_kept_entry_id: "kept".into(), + messages_to_summarize: history.clone(), + turn_prefix_messages: vec![user("the start of a long task")], + tokens_before: 100, + previous_summary: Some("## Goal\nthe old goal".into()), + file_operations: FileOperations::default(), + // A small reserve no longer limits the summary. + settings: CompactionSettings { + reserve_tokens: 100, + ..settings(20) + }, + }; + let faux = FauxProvider::new(); + if !history.is_empty() { + faux.push_text("## Goal\nhistory"); + } + faux.push_text("## Original Request\ntask"); + let model = faux.model(); + let summarizer = Summarizer { + model: &model, + stream_fn: &faux, + // The session's own output limit does not carry over either. + options: StreamOptions { + max_tokens: Some(64), + ..StreamOptions::default() + }, + max_retries: 0, + retry_base_ms: 1, + retry_max_ms: 60_000, + on_retry: None, + }; + compact(&preparation, &summarizer, Some("keep paths")) + .await + .unwrap(); + let requests = faux.requests(); + assert_eq!(requests.len(), 1 + usize::from(!history.is_empty())); + assert!( + requests + .iter() + .all(|request| request.options.max_tokens.is_none()) + ); + let prompt = last_prompt(&faux); + assert!(prompt.starts_with("# Conversation\n[User]: the start of a long task")); + assert!( + prompt.ends_with("\n\nAdditional focus: keep paths"), + "{prompt}" + ); + } + } + #[tokio::test] async fn a_branch_summary_sees_the_branch_as_the_context_did() { let mut session = SessionManager::in_memory("/work"); @@ -1367,26 +1445,8 @@ mod tests { "{prompt}" ); assert!(!prompt.contains("[Tool result]") && !prompt.contains("fn parse")); - assert_eq!(faux.requests()[0].options.max_tokens, Some(4_096)); - - // The output cap is the model's when that is lower. - let mut small = FauxProvider::default_model(); - small.max_tokens = 1_024; - let faux = FauxProvider::with_model(small.clone()); - faux.push_text("## Goal\nFix the parser"); - let summarizer = Summarizer { - model: &small, - stream_fn: &faux, - options: StreamOptions::default(), - max_retries: 0, - retry_base_ms: 1, - retry_max_ms: 60_000, - on_retry: None, - }; - summarize_branch(&session.branch(), &summarizer, &settings(1), 0, None, false) - .await - .unwrap(); - assert_eq!(faux.requests()[0].options.max_tokens, Some(1_024)); + // No fixed output cap: the provider fits the model's own to the request. + assert_eq!(faux.requests()[0].options.max_tokens, None); } #[tokio::test] diff --git a/apps/maple-agent/crates/pi-coding-agent/src/extensions.rs b/apps/maple-agent/crates/pi-coding-agent/src/extensions.rs index ecb22c959..a227a5f10 100644 --- a/apps/maple-agent/crates/pi-coding-agent/src/extensions.rs +++ b/apps/maple-agent/crates/pi-coding-agent/src/extensions.rs @@ -37,7 +37,7 @@ use tokio_util::sync::CancellationToken; use crate::agent_session::{CompactionReason, SessionCore, StreamingBehavior}; use crate::bash_executor::BashResult; use crate::compaction::{CompactionPreparation, CompactionResult}; -use crate::messages::SessionMessage; +use crate::messages::{CustomMessage, SessionMessage}; use crate::session::{ProjectedEntry, SessionEntry}; use crate::system_prompt::SystemPromptOptions; use crate::tools::BashOperations; @@ -285,16 +285,28 @@ pub enum InputAction { Handled, } -/// A prompt about to start a run. +/// A run about to start, from a prompt or from a custom message, as Pi's +/// `before_agent_start` for both of its message types. #[derive(Clone, Debug)] pub struct BeforeAgentStart { - pub prompt: String, - /// The images attached to the prompt. - pub images: Vec, + /// The message that starts the run. + pub trigger: RunTrigger, pub system_prompt: String, pub options: SystemPromptOptions, } +/// What starts a run, as Pi's `messageType` says. +#[derive(Clone, Debug, PartialEq)] +pub enum RunTrigger { + /// A prompt: its text after expansion, and the images attached to it. + User { + prompt: String, + images: Vec, + }, + /// A custom message that starts a turn, as Pi's `sendMessage` with `triggerTurn`. + Custom(CustomMessage), +} + /// A message an extension asks to add. #[derive(Clone, Debug, PartialEq)] pub struct CustomMessageDraft { diff --git a/apps/maple-agent/crates/pi-coding-agent/src/resources.rs b/apps/maple-agent/crates/pi-coding-agent/src/resources.rs index 91fa61032..e231aca1b 100644 --- a/apps/maple-agent/crates/pi-coding-agent/src/resources.rs +++ b/apps/maple-agent/crates/pi-coding-agent/src/resources.rs @@ -123,8 +123,17 @@ pub fn load_context_files(cwd: &Path, agent_dir: &Path) -> Vec { let shadowed = shadowed_context_file(cwd); let mut ancestors: Vec = cwd .ancestors() - .filter_map(context_file_in) - .filter(|file| shadowed.as_ref() != Some(&canonical(&file.path))) + .filter_map(|dir| { + let file = context_file_in(dir)?; + // The folder's real path names the file, not the file's own: a worktree's + // file that links to the main checkout's resolves to the shadowed one. + let is_shadowed = shadowed.as_ref().is_some_and(|shadowed| { + file.path + .file_name() + .is_some_and(|name| canonical(dir).join(name) == *shadowed) + }); + (!is_shadowed).then_some(file) + }) .collect(); ancestors.reverse(); for file in ancestors { @@ -877,21 +886,30 @@ pub fn parse_skill_block(text: &str) -> Option { pub fn parse_command_args(input: &str) -> Vec { let mut args = Vec::new(); let mut current = String::new(); + // Whether an argument has started, so empty quotes (`""` or `''`) keep their place, + // as in Pi. + let mut started = false; let mut quote = None; for ch in input.chars() { match quote { Some(open) if ch == open => quote = None, Some(_) => current.push(ch), - None if ch == '"' || ch == '\'' => quote = Some(ch), + None if ch == '"' || ch == '\'' => { + quote = Some(ch); + started = true; + } None if ch.is_whitespace() => { - if !current.is_empty() { + if std::mem::take(&mut started) { args.push(std::mem::take(&mut current)); } } - None => current.push(ch), + None => { + current.push(ch); + started = true; + } } } - if !current.is_empty() { + if started { args.push(current); } args @@ -1177,6 +1195,16 @@ mod tests { assert_eq!(contents(&cwd), ["main", "worktree"]); // The main checkout itself loads its own. assert_eq!(contents(&repo), ["main"]); + // A worktree file that links to the main checkout's is the worktree's own. + #[cfg(unix)] + { + fs::remove_file(worktree.join("CLAUDE.md")).unwrap(); + std::os::unix::fs::symlink(repo.join("AGENTS.md"), worktree.join("AGENTS.md")).unwrap(); + let files = load_context_files(&cwd, &paths.agent_dir); + assert_eq!(files.len(), 1); + assert_eq!(files[0].path, worktree.join("AGENTS.md")); + assert_eq!(files[0].content, "main"); + } } #[test] @@ -1361,6 +1389,21 @@ mod tests { assert_eq!(substitute_args("${@:99999999999999999999}", &args), ""); } + #[test] + fn empty_quoted_arguments_keep_their_place() { + assert_eq!(parse_command_args(r#""" " ""#), ["", " "]); + assert_eq!(parse_command_args("a '' c"), ["a", "", "c"]); + assert_eq!(parse_command_args(r#"a """#), ["a", ""]); + // So a positional default applies and later arguments do not shift. + let template = "Review $2. Focus: ${1:-correctness}."; + for args in [r#""" src/index.ts"#, "'' src/index.ts"] { + assert_eq!( + substitute_args(template, &parse_command_args(args)), + "Review src/index.ts. Focus: correctness." + ); + } + } + #[test] fn templates_load_and_expand_by_name() { let (_root, cwd, paths) = layout(); diff --git a/apps/maple-agent/crates/pi-coding-agent/src/tools/shell.rs b/apps/maple-agent/crates/pi-coding-agent/src/tools/shell.rs index 0cb7856cc..2c0d100e0 100644 --- a/apps/maple-agent/crates/pi-coding-agent/src/tools/shell.rs +++ b/apps/maple-agent/crates/pi-coding-agent/src/tools/shell.rs @@ -40,10 +40,36 @@ impl ShellConfig { } fn is_legacy_wsl_bash(path: &Path) -> bool { - let normalized = path - .to_string_lossy() - .replace('/', "\\") - .to_ascii_lowercase(); + is_legacy_wsl_bash_in(path, system_root().as_deref()) +} + +/// Whether `path` is the WSL launcher in the Windows folder `system_root`, or, when that +/// is unknown, in the conventional one on any drive. +fn is_legacy_wsl_bash_in(path: &Path, system_root: Option<&str>) -> bool { + let normalize = |path: &str| -> String { + let mut normalized = String::with_capacity(path.len()); + for ch in path.chars() { + let ch = if ch == '/' { '\\' } else { ch }; + // Repeated separators are one, as Node's `win32.normalize` reads them. + if !(ch == '\\' && normalized.ends_with('\\')) { + normalized.push(ch.to_ascii_lowercase()); + } + } + normalized + }; + let normalized = normalize(&path.to_string_lossy()); + if let Some(root) = system_root { + let root = normalize(root); + let bytes = root.as_bytes(); + // Only a drive path names the Windows folder. + if !(bytes.len() >= 3 && bytes[0].is_ascii_alphabetic() && &bytes[1..3] == b":\\") { + return false; + } + let root = root.trim_end_matches('\\'); + return ["system32", "sysnative"] + .iter() + .any(|dir| normalized == format!("{root}\\{dir}\\bash.exe")); + } let Some((drive, rest)) = normalized.split_once(':') else { return false; }; @@ -55,19 +81,37 @@ fn is_legacy_wsl_bash(path: &Path) -> bool { ) } -/// The first file named `name` on PATH. +/// The Windows folder, when the environment names one. +fn system_root() -> Option { + std::env::var("SystemRoot") + .ok() + .filter(|root| !root.is_empty()) +} + +/// Whether `path` is a Microsoft Store alias: a link that runs the app, which cannot be +/// opened through as an ordinary file or link can. Only Windows has them. +fn is_store_alias(path: &Path) -> bool { + cfg!(windows) && std::fs::metadata(path).is_err() && std::fs::symlink_metadata(path).is_ok() +} + +/// Whether a shell exists at `path`, as Pi checks it, Microsoft Store aliases included. +fn shell_path_exists(path: &Path) -> bool { + path.exists() || is_store_alias(path) +} + +/// The first file named `name` on PATH, a Microsoft Store alias included. pub(crate) fn find_executable_on_path(name: &str) -> Option { let path = std::env::var_os("PATH")?; std::env::split_paths(&path) .map(|dir| dir.join(name)) - .find(|candidate| candidate.is_file()) + .find(|candidate| candidate.is_file() || is_store_alias(candidate)) } /// The shell for `bash`: `custom_path` when given, else on Windows Git Bash, then /// `bash.exe` on PATH; elsewhere `/bin/bash`, then `bash` on PATH, then `sh`. pub fn shell_config(custom_path: Option<&Path>) -> Result { if let Some(path) = custom_path { - if path.exists() { + if shell_path_exists(path) { return Ok(ShellConfig::bash(path.to_path_buf())); } return Err(format!("Custom shell path not found: {}", path.display())); @@ -79,7 +123,7 @@ pub fn shell_config(custom_path: Option<&Path>) -> Result { searched.push(PathBuf::from(dir).join("Git").join("bin").join("bash.exe")); } } - if let Some(found) = searched.iter().find(|path| path.exists()) { + if let Some(found) = searched.iter().find(|path| shell_path_exists(path)) { return Ok(ShellConfig::bash(found.clone())); } if let Some(found) = find_executable_on_path("bash.exe") { @@ -95,7 +139,7 @@ pub fn shell_config(custom_path: Option<&Path>) -> Result { )); } let system_bash = Path::new("/bin/bash"); - if system_bash.exists() { + if shell_path_exists(system_bash) { return Ok(ShellConfig::bash(system_bash.to_path_buf())); } if let Some(found) = find_executable_on_path("bash") { @@ -116,13 +160,27 @@ pub const POWERSHELL_ARGS: [&str; 5] = [ "-Command", ]; -/// PowerShell on Windows, PowerShell 7 when it is installed. +/// PowerShell on Windows, PowerShell 7 when it is installed. Each is also looked for in +/// its installation folder, which need not be on PATH or on drive C. pub fn powershell_config() -> Result { if !cfg!(windows) { return Err("The powershell tool is only available on Windows.".to_string()); } + let installed = |dir: Option, path: &[&str]| { + dir.map(|dir| path.iter().fold(dir, |dir, part| dir.join(part))) + .filter(|path| shell_path_exists(path)) + }; + let program_files = std::env::var_os("ProgramFiles").map(PathBuf::from); + let system_root = system_root().map(PathBuf::from); let shell = find_executable_on_path("pwsh.exe") + .or_else(|| installed(program_files, &["PowerShell", "7", "pwsh.exe"])) .or_else(|| find_executable_on_path("powershell.exe")) + .or_else(|| { + installed( + system_root, + &["System32", "WindowsPowerShell", "v1.0", "powershell.exe"], + ) + }) .ok_or_else(|| { "No PowerShell executable found. Install PowerShell or add powershell.exe/pwsh.exe to PATH." .to_string() @@ -213,7 +271,7 @@ pub fn kill_process_tree(pid: u32) { use std::os::windows::process::CommandExt; const CREATE_NO_WINDOW: u32 = 0x0800_0000; // The trusted System32 copy, so cleanup does not depend on PATH. - let system_root = std::env::var_os("SystemRoot").unwrap_or_else(|| "C:\\Windows".into()); + let system_root = system_root().unwrap_or_else(|| "C:\\Windows".into()); let taskkill = PathBuf::from(system_root) .join("System32") .join("taskkill.exe"); @@ -241,6 +299,27 @@ mod tests { assert_eq!(git_bash.transport, CommandTransport::Argv); } + #[test] + fn the_wsl_launcher_is_found_in_the_windows_folder() { + let launcher = + |path: &str, root: Option<&str>| is_legacy_wsl_bash_in(Path::new(path), root); + // Wherever Windows is installed, when the environment says where. + assert!(launcher("D:\\Win\\System32\\bash.exe", Some("D:\\Win"))); + assert!(launcher("d:/win//sysnative/BASH.EXE", Some("D:\\Win\\"))); + assert!(!launcher( + "C:\\Windows\\System32\\bash.exe", + Some("D:\\Win") + )); + // A Windows folder that is no drive path names no launcher. + assert!(!launcher( + "C:\\Windows\\System32\\bash.exe", + Some("Windows") + )); + // Otherwise the conventional folder on any drive. + assert!(launcher("E:\\Windows\\System32\\bash.exe", None)); + assert!(!launcher("C:\\Program Files\\Git\\bin\\bash.exe", None)); + } + #[test] fn a_missing_custom_shell_is_an_error() { assert_eq!( diff --git a/apps/maple-agent/crates/pi-coding-agent/tests/agent_session.rs b/apps/maple-agent/crates/pi-coding-agent/tests/agent_session.rs index 93ddce4a3..4b1399655 100644 --- a/apps/maple-agent/crates/pi-coding-agent/tests/agent_session.rs +++ b/apps/maple-agent/crates/pi-coding-agent/tests/agent_session.rs @@ -11,10 +11,11 @@ use pi_agent_core::QueueMode; use pi_agent_core::{AgentToolResult, FnTool}; use pi_ai::faux::{FauxProvider, faux_message}; use pi_ai::{AssistantContent, Content, ImageContent, Message, StopReason, ThinkingLevel, Usage}; -use pi_coding_agent::compaction::estimate_tokens; +use pi_coding_agent::compaction::{CompactionResult, estimate_tokens}; use pi_coding_agent::extensions::{ - BeforeForkResult, BeforeTreeResult, CustomMessageDraft, ForkPosition, Input, InputAction, - ModelSelect, SessionBeforeFork, SessionBeforeTree, SlashCommandSource, extension, + BeforeCompactResult, BeforeForkResult, BeforeTreeResult, CustomMessageDraft, ForkPosition, + Input, InputAction, ModelSelect, SessionBeforeCompact, SessionBeforeFork, SessionBeforeTree, + SlashCommandSource, extension, }; use pi_coding_agent::resources::{PromptTemplate, ResourceSource, Resources, Skill}; use pi_coding_agent::session::{EntryKind, SessionEntry, SessionHeader, SessionManager}; @@ -1500,6 +1501,104 @@ async fn a_reply_cut_off_by_a_full_context_is_compacted_and_asked_again() { ); } +#[tokio::test] +async fn recovering_a_cut_off_tool_call_omits_its_result_with_it() { + // A resumed session ends with a reply cut off at the output limit, and a result for + // its call saying it was not run (Pi #11227). + let harness = Harness::new(); + let model = harness.faux.model(); + let mut tree = SessionManager::in_memory("/work"); + tree.append_model_change(&model.provider, &model.id); + tree.append_message(SessionMessage::Llm(Message::User( + pi_ai::UserMessage::text("check status"), + ))); + let mut cut_off = faux_message(vec![AssistantContent::tool_call( + "call_cut_off", + "bash", + json!({ "command": "echo" }), + )]); + cut_off.stop_reason = StopReason::Length; + cut_off.usage = Usage { + input: 100, + output: 16, + total_tokens: 116, + ..Usage::default() + }; + tree.append_message(SessionMessage::Llm(Message::Assistant(cut_off))); + tree.append_message(SessionMessage::Llm(Message::ToolResult( + pi_ai::ToolResultMessage { + tool_call_id: "call_cut_off".into(), + tool_name: "bash".into(), + content: vec![Content::text("Not executed: output token limit reached.")], + details: None, + usage: None, + is_error: true, + timestamp: 0, + }, + ))); + let supply = extension("supply", |api| { + api.on(|event: SessionBeforeCompact, _ctx| async move { + Ok(BeforeCompactResult { + cancel: false, + compaction: Some(CompactionResult { + summary: "summary".into(), + first_kept_entry_id: event.preparation.first_kept_entry_id, + tokens_before: event.preparation.tokens_before, + estimated_tokens_after: None, + usage: Default::default(), + details: None, + }), + }) + }); + }); + let session = harness + .session_with(|options| { + options.session = tree; + options.extensions = vec![supply]; + }) + .await; + harness.faux.push_text("done"); + let events = record(&session); + + session + .prompt("continue", PromptOptions::default()) + .await + .unwrap(); + let requests = harness.faux.requests(); + let sent = &requests.last().unwrap().context.messages; + assert!( + sent.iter().all(|message| match message { + Message::ToolResult(_) => false, + Message::Assistant(reply) => reply.stop_reason != StopReason::Length, + _ => true, + }), + "{sent:?}" + ); + // Both are omitted by entries of their own. + let omitted = session.with_session(|tree| { + tree.entries() + .iter() + .filter(|entry| { + matches!( + entry.kind, + EntryKind::ContextEdit { + replacement: None, + .. + } + ) + }) + .count() + }); + assert_eq!(omitted, 2); + let appended = events + .lock() + .unwrap() + .iter() + .filter(|event| event.as_str() == "entry:context_edit") + .count(); + assert_eq!(appended, 2); +} + /// A store whose every write fails. struct ReadOnlyStore; diff --git a/apps/maple-agent/crates/pi-coding-agent/tests/extensions.rs b/apps/maple-agent/crates/pi-coding-agent/tests/extensions.rs index 2ede8058b..d3930c25e 100644 --- a/apps/maple-agent/crates/pi-coding-agent/tests/extensions.rs +++ b/apps/maple-agent/crates/pi-coding-agent/tests/extensions.rs @@ -13,9 +13,9 @@ use pi_coding_agent::compaction::CompactionResult; use pi_coding_agent::extensions::{ AfterProviderResponse, AgentEventSeen, AgentSettled, BeforeAgentStart, BeforeAgentStartResult, BeforeCompactResult, BeforeProviderHeaders, Context, ContextWithSystem, CustomMessageDraft, - Input, InputAction, MessageEnd, ProviderStreamEvent, SessionBeforeCompact, SessionCompact, - SessionCompactFailed, SessionInfoChanged, SessionStart, SlashCommandInfo, SlashCommandSource, - ToolCall, ToolPrompt, ToolResult, extension, + Input, InputAction, MessageEnd, ProviderStreamEvent, RunTrigger, SessionBeforeCompact, + SessionCompact, SessionCompactFailed, SessionInfoChanged, SessionStart, SlashCommandInfo, + SlashCommandSource, ToolCall, ToolPrompt, ToolResult, extension, }; use pi_coding_agent::session::{EntryKind, SessionManager}; use pi_coding_agent::store::JsonlStore; @@ -293,8 +293,9 @@ async fn before_agent_start_adds_messages_and_can_force_the_prompt_for_one_run() display: false, details: None, }), - system_prompt: (event.prompt == "forced") - .then(|| "You only answer in haiku.".to_string()), + system_prompt: + matches!(&event.trigger, RunTrigger::User { prompt, .. } if prompt == "forced") + .then(|| "You only answer in haiku.".to_string()), options: None, }) }); @@ -975,6 +976,207 @@ async fn extension_messages_queued_during_a_run_are_not_listed() { run.await.unwrap().unwrap(); } +fn custom_draft(custom_type: &str, text: &str) -> CustomMessageDraft { + CustomMessageDraft { + custom_type: custom_type.into(), + content: vec![Content::text(text)], + display: true, + details: None, + } +} + +/// The text of every user message of request `index`, in order. +fn user_texts(harness: &Harness, index: usize) -> Vec { + harness.faux.requests()[index] + .context + .messages + .iter() + .filter_map(|message| match message { + Message::User(user) => Some(content_text(&user.content)), + _ => None, + }) + .collect() +} + +#[tokio::test] +async fn a_run_a_custom_message_starts_keeps_the_prompt_for_every_request() { + // Pi #10739: such a run skipped `before_agent_start`, so its next request patched away + // the sections the handlers had added. + let harness = Harness::new(); + harness.faux.push_text("hi"); + harness.faux.push_tool_call("echo", json!({ "text": "a" })); + harness.faux.push_text("done"); + let triggers = Arc::new(Mutex::new(Vec::new())); + let seen = triggers.clone(); + let rules = extension("rules", move |api| { + let seen = seen.clone(); + api.on(move |event: BeforeAgentStart, _ctx| { + seen.lock().unwrap().push(event.trigger.clone()); + async move { + let mut options = event.options; + options + .sections + .insert("extension_rules".into(), "Answer in one word.".into()); + Ok(BeforeAgentStartResult { + options: Some(options), + ..BeforeAgentStartResult::default() + }) + } + }); + }); + let session = harness + .session_with(|options| { + options.tools = vec![echo_tool(true)]; + options.extensions = vec![rules]; + }) + .await; + session + .prompt("hello", PromptOptions::default()) + .await + .unwrap(); + session + .send_custom_message(custom_draft("task-done", "background task finished")) + .await + .unwrap(); + + let prompts: Vec = harness + .faux + .requests() + .iter() + .map(|request| pi_ai::transcript::current_system_prompt(&request.context.messages)) + .collect(); + assert_eq!(prompts.len(), 3); + assert!(prompts[0].contains("Answer in one word.")); + assert!(prompts.iter().all(|prompt| *prompt == prompts[0])); + let triggers = triggers.lock().unwrap(); + assert!(matches!(&triggers[0], RunTrigger::User { prompt, .. } if prompt == "hello")); + assert!( + matches!(&triggers[1], RunTrigger::Custom(message) if message.custom_type == "task-done") + ); +} + +/// An extension whose `before_agent_start` handler waits for `release` when `blocks` +/// says so, after telling `started`. +fn blocking_start( + blocks: fn(&RunTrigger) -> bool, + started: Arc, + release: Arc, +) -> Arc { + extension("blocking", move |api| { + let (started, release) = (started.clone(), release.clone()); + api.on(move |event: BeforeAgentStart, _ctx| { + let block = blocks(&event.trigger); + let (started, release) = (started.clone(), release.clone()); + async move { + if block { + started.notify_one(); + release.notified().await; + } + Ok(BeforeAgentStartResult::default()) + } + }); + }) +} + +#[tokio::test] +async fn a_custom_message_sent_while_a_run_starts_joins_it() { + let harness = Harness::new(); + harness.faux.push_text("handled both"); + let (started, release) = ( + Arc::new(tokio::sync::Notify::new()), + Arc::new(tokio::sync::Notify::new()), + ); + let blocking = blocking_start( + |trigger| matches!(trigger, RunTrigger::Custom(message) if message.custom_type == "first"), + started.clone(), + release.clone(), + ); + let session = harness.session_with_extensions(vec![blocking]).await; + let runner = session.clone(); + let first = tokio::spawn(async move { + runner + .send_custom_message(custom_draft("first", "task one finished")) + .await + }); + started.notified().await; + assert!(session.is_streaming()); + assert_eq!( + session + .send_custom_message(custom_draft("second", "task two finished")) + .await + .unwrap(), + PromptOutcome::Queued + ); + release.notify_one(); + first.await.unwrap().unwrap(); + + assert_eq!(harness.faux.requests().len(), 1); + assert_eq!( + user_texts(&harness, 0), + ["task one finished", "task two finished"] + ); +} + +#[tokio::test] +async fn a_message_added_while_a_run_starts_is_in_its_first_request() { + let harness = Harness::new(); + harness.faux.push_text("done"); + let note = extension("note", |api| { + api.on(|_event: BeforeAgentStart, ctx| { + ctx.send_message( + custom_draft("note", "context from handler"), + Delivery::Append, + ); + async { Ok(BeforeAgentStartResult::default()) } + }); + }); + let session = harness.session_with_extensions(vec![note]).await; + session + .prompt("hello", PromptOptions::default()) + .await + .unwrap(); + + assert_eq!(harness.faux.requests().len(), 1); + // It goes before the prompt, as it would have while idle. + assert_eq!(user_texts(&harness, 0), ["context from handler", "hello"]); +} + +#[tokio::test] +async fn an_abort_while_a_run_starts_stops_it_before_its_first_request() { + let harness = Harness::new(); + harness.faux.push_text("should not complete"); + let (started, release) = ( + Arc::new(tokio::sync::Notify::new()), + Arc::new(tokio::sync::Notify::new()), + ); + let blocking = blocking_start(|_| true, started.clone(), release.clone()); + let session = harness.session_with_extensions(vec![blocking]).await; + let events = record(&session); + let runner = session.clone(); + let run = tokio::spawn(async move { + runner + .send_custom_message(custom_draft("task-done", "background task finished")) + .await + }); + started.notified().await; + session.abort(); + release.notify_one(); + run.await.unwrap().unwrap(); + + let last = session.messages().pop().unwrap(); + assert!(matches!( + last, + SessionMessage::Llm(Message::Assistant(reply)) if reply.stop_reason == pi_ai::StopReason::Aborted + )); + assert!(!session.is_streaming()); + assert!( + events + .lock() + .unwrap() + .contains(&"settled:aborted".to_string()) + ); +} + fn provider_extension(provider: FauxProvider) -> Arc { extension("provider", move |api| { api.register_provider( @@ -1029,7 +1231,8 @@ async fn prompt_options_from_before_agent_start_last_for_one_run() { harness.faux.push_text("two"); let notes = extension("notes", |api| { api.on(|event: BeforeAgentStart, _ctx| async move { - if event.prompt != "with notes" { + if !matches!(&event.trigger, RunTrigger::User { prompt, .. } if prompt == "with notes") + { return Ok(BeforeAgentStartResult::default()); } let mut options = event.options; diff --git a/apps/maple-agent/crates/store-bench/src/data.rs b/apps/maple-agent/crates/store-bench/src/data.rs index 69811784e..d3837eee2 100644 --- a/apps/maple-agent/crates/store-bench/src/data.rs +++ b/apps/maple-agent/crates/store-bench/src/data.rs @@ -205,6 +205,7 @@ pub fn entry(task: &TaskSpec, seq: usize) -> SessionEntry { id: format!("call-{seq}"), name: (if seq % 61 == 1 { "todo_write" } else { "read" }).into(), arguments: arguments.as_object().cloned().unwrap_or_else(Map::new), + arguments_error: None, partial_json: None, })); } @@ -241,6 +242,7 @@ pub fn oversized_entry(task: &TaskSpec, seq: usize, len: usize) -> SessionEntry id: format!("oversized-{seq}"), name: "write".into(), arguments: arguments.as_object().unwrap().clone(), + arguments_error: None, partial_json: None, })], api: "openai-completions".into(),