Skip to content

fix(a2a)!: admit messages into the interactive session #269

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
danielkov wants to merge 3 commits into
base: main
Choose a base branch
from
Open

fix(a2a)!: admit messages into the interactive session #269

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
1c77932
fix(a2a)!: admit messages into the interactive session
danielkov Oct 13, 2026
766497c
Merge remote-tracking branch 'origin/main' into fix/a2a-session-admis…
danielkov Oct 13, 2026
7ea1614
test(a2a): adapt prompt preparation to current main
danielkov Oct 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 33 additions & 1 deletion src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1964,6 +1964,17 @@ async fn termination_signal() -> io::Result<()> {
}
}

fn serve_a2a_execution_mode(
no_stdio: bool,
protocol: AcpProtocolVersion,
) -> kit::protocols::http::A2aExecutionMode {
if !no_stdio && matches!(protocol, AcpProtocolVersion::V2) {
kit::protocols::http::A2aExecutionMode::SessionAdmission
} else {
kit::protocols::http::A2aExecutionMode::Standalone
}
}

async fn supervise_serve(
runtime: std::sync::Arc<kit::Runtime>,
sessions: kit::protocols::acp::SessionRegistry,
Expand Down Expand Up @@ -2299,13 +2310,14 @@ async fn run_cli(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
let address = a2a.unwrap_or_else(|| "127.0.0.1:0".into());
let serve_a2a = !no_a2a;
let sessions = kit::protocols::acp::SessionRegistry::new();
let http = kit::protocols::http::start_with_registry(
let http = kit::protocols::http::start_with_registry_and_mode(
runtime.clone(),
address,
serve_a2a,
remote_acp,
server_credential_file.as_deref(),
sessions.clone(),
serve_a2a_execution_mode(no_stdio, stdio_protocol_version),
)
.await?;
let bound = http.address();
Expand Down Expand Up @@ -2600,6 +2612,26 @@ mod tests {
supervise_serve_with_trigger, validate_auth_storage,
};

#[test]
fn serve_a2a_execution_mode_preserves_no_stdio_and_v1_execution() {
use super::{AcpProtocolVersion, serve_a2a_execution_mode};
use kit::protocols::http::A2aExecutionMode;
assert_eq!(
serve_a2a_execution_mode(false, AcpProtocolVersion::V2),
A2aExecutionMode::SessionAdmission
);
for (no_stdio, protocol) in [
(true, AcpProtocolVersion::V1),
(true, AcpProtocolVersion::V2),
(false, AcpProtocolVersion::V1),
] {
assert_eq!(
serve_a2a_execution_mode(no_stdio, protocol),
A2aExecutionMode::Standalone
);
}
}

#[test]
fn usage_command_accepts_optional_provider_and_credential_overrides() {
for provider in [
Expand Down
269 changes: 261 additions & 8 deletions src/protocols/a2a.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use a2a_protocol_server::{
use a2a_protocol_types::{
agent_card::{AgentCapabilities, AgentCard, AgentInterface, AgentSkill},
error::A2aResult,
message::Part,
message::{Message, MessageId, MessageRole, Part},
security::{HttpAuthSecurityScheme, SecurityRequirement, SecurityScheme, StringList},
task::TaskState,
};
Expand All @@ -17,7 +17,167 @@ use tracing::Instrument as _;

use crate::runtime::Runtime;

struct KitAgent(Arc<Runtime>);
pub(crate) struct Dispatcher {
inner: JsonRpcDispatcher,
sessions: Option<crate::protocols::acp::SessionRegistry>,
}

impl Dispatcher {
pub(crate) async fn dispatch(
&self,
request: axum::http::Request<hyper::body::Incoming>,
) -> axum::http::Response<axum::body::Body> {
use axum::{
body::{Body, to_bytes},
http::{Method, Response, StatusCode},
};
let Some(sessions) = &self.sessions else {
return self.inner.dispatch(request).await.map(Body::new);
};
if request.method() != Method::POST {
return self.inner.dispatch(request).await.map(Body::new);
}
let bytes = match to_bytes(Body::new(request.into_body()), 4 * 1024 * 1024).await {

@kit-code-agent kit-code-agent Bot Oct 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] Preserve the HTTP media-type boundary before admitting messages

With session admission enabled, every POST reaches this body parser without the Content-Type or A2A-Version checks performed by JsonRpcDispatcher. On a listener without a configured bearer token, a hostile web page can therefore send a valid admission request as a CORS-safelisted text/plain POST, without the custom version header, and inject a prompt into the bound ACP session in browsers that permit the loopback connection. Blocking response access does not prevent admission, and the shared listener does not check Origin. The pinned dispatcher previously rejected this request before executing it. Preserve the JSON media-type and version validation before admission, and add a regression test ensuring a text/plain request cannot enqueue a prompt or fall back to standalone execution.

For agents: 
Validate the following issue, address if needed:

<comment>Session-admission POSTs bypass HTTP media-type and version validation, allowing browser-simple text/plain requests to enqueue executable prompts on an unauthenticated listener.</comment>
<file_context>
+        let Some(sessions) = &self.sessions else {
+            return self.inner.dispatch(request).await.map(Body::new);
+        };
+        if request.method() != Method::POST {
+            return self.inner.dispatch(request).await.map(Body::new);
+        }
+        let bytes = match to_bytes(Body::new(request.into_body()), 4 * 1024 * 1024).await {
+            Ok(bytes) => bytes,
+            Err(_) => {
+                let mut response =
+                    Response::new(Body::from("A2A request body exceeds admission limit"));
+                *response.status_mut() = StatusCode::PAYLOAD_TOO_LARGE;
+                return response;
+            }
+        };
</file_context>

@kit-code-agent kit-code-agent Bot Oct 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Retain the request-body deadline in session-admission mode

The direct to_bytes(...).await retains the 4 MiB size limit but drops the pinned JsonRpcDispatcher's default 30-second body-read timeout. A client can send POST headers and an unfinished body below the cap, then leave this future pending indefinitely. The shared listener spawns a connection task for each accepted socket and supplies no replacement body deadline, so repeated incomplete requests retain tasks, sockets, and any accumulated body buffers until the clients disconnect or the server shuts down. Bound body collection with a deadline equivalent to the existing dispatcher, return a distinct timeout error without admission or standalone fallback, and cover an incomplete request body in a regression test.

For agents: 
Validate the following issue, address if needed:

<comment>The session-admission body reader removes the dispatcher's 30-second deadline, allowing incomplete POST bodies to retain connection resources indefinitely.</comment>
<file_context>
+        if request.method() != Method::POST {
+            return self.inner.dispatch(request).await.map(Body::new);
+        }
+        let bytes = match to_bytes(Body::new(request.into_body()), 4 * 1024 * 1024).await {
+            Ok(bytes) => bytes,
+            Err(_) => {
+                let mut response =
+                    Response::new(Body::from("A2A request body exceeds admission limit"));
+                *response.status_mut() = StatusCode::PAYLOAD_TOO_LARGE;
+                return response;
+            }
+        };
</file_context>

Ok(bytes) => bytes,
Err(_) => {
let mut response =
Response::new(Body::from("A2A request body exceeds admission limit"));
*response.status_mut() = StatusCode::PAYLOAD_TOO_LARGE;
return response;
}
};
let value = match serde_json::from_slice::<serde_json::Value>(&bytes) {
Ok(value) => value,
Err(_) => {
return rpc_response(
serde_json::Value::Null,
Err((-32700, "Invalid JSON".into())),
);
}
};
let id = value
.get("id")
.filter(|id| id.is_null() || id.is_string() || id.is_number())
.cloned()
.unwrap_or_default();
rpc_response(id, admit_message(sessions, value).await)
}
}

fn rpc_response(
id: serde_json::Value,
result: Result<Message, (i32, String)>,
) -> axum::http::Response<axum::body::Body> {
use serde_json::{Map, Value};
let result = result.and_then(|message| {
serde_json::to_value(a2a_protocol_types::responses::SendMessageResponse::Message(
message,
))
.map_err(|error| (-32603, error.to_string()))
});
let mut value = Map::from_iter([
("jsonrpc".into(), Value::String("2.0".into())),
("id".into(), id),
]);
match result {
Ok(message) => {
value.insert("result".into(), message);
}
Err((code, message)) => {
value.insert(
"error".into(),
Value::Object(Map::from_iter([
("code".into(), Value::from(code)),
("message".into(), Value::String(message)),
])),
);
}
}
let value = Value::Object(value);
let mut response = axum::http::Response::new(axum::body::Body::from(value.to_string()));
response.headers_mut().insert(
axum::http::header::CONTENT_TYPE,
axum::http::HeaderValue::from_static("application/json"),
);
response
}

async fn admit_message(
sessions: &crate::protocols::acp::SessionRegistry,
request: serde_json::Value,
) -> Result<Message, (i32, String)> {
if request["jsonrpc"] != "2.0"
|| !request.is_object()
|| !request
.get("id")
.is_some_and(|id| id.is_null() || id.is_string() || id.is_number())
{
return Err((-32600, "Invalid JSON-RPC request".into()));
}
let Some(method) = request["method"].as_str() else {
return Err((-32600, "JSON-RPC method must be a string".into()));
};
if !matches!(method, "SendMessage" | "message/send") {
return Err((-32601, "Session-bound A2A supports SendMessage admission acknowledgments, not task tracking or streaming".into()));
}
let params: a2a_protocol_types::params::MessageSendParams =
serde_json::from_value(request["params"].clone())
.map_err(|error| (-32602, error.to_string()))?;
if params.configuration.as_ref().is_some_and(|configuration| {
configuration.return_immediately == Some(true)
|| configuration.task_push_notification_config.is_some()
}) {
return Err((-32602, "Session-bound A2A requires a direct acknowledgment: returnImmediately=false and no push notification configuration".into()));
}
if params.message.role != MessageRole::User || params.message.task_id.is_some() {
return Err((
-32602,
"Admission requires a new user message without a taskId".into(),
));
}
if params
.message
.parts
.iter()
.any(|part| part.text_content().is_none())
{
return Err((
-32602,
"Session admission supports text parts only; nontext parts are not accepted".into(),
));
}
let prompt = params
.message
.parts
.iter()
.filter_map(Part::text_content)
.collect::<Vec<_>>()
.join("\n");
if prompt.trim().is_empty() {
return Err((-32602, "A2A request must contain a text part".into()));
}
let acknowledgment = sessions
.admit_a2a(prompt)
.await
.map_err(|error| (-32603, error))?;
Ok(Message {
id: MessageId::new(format!("{}-admission", params.message.id)),
role: MessageRole::Agent,
parts: vec![Part::text(acknowledgment)],
task_id: None,
context_id: params.message.context_id,
reference_task_ids: None,
extensions: None,
metadata: Some(serde_json::Value::Object(serde_json::Map::from_iter([(
"kind".into(),
serde_json::Value::String("admission_acknowledgment".into()),
)]))),
})
}

struct KitAgent {
runtime: Arc<Runtime>,
}

impl AgentExecutor for KitAgent {
fn execute<'a>(
Expand All @@ -28,7 +188,7 @@ impl AgentExecutor for KitAgent {
Box::pin(
async move {
let emit = EventEmitter::new(context, queue);
emit.status(TaskState::Working).await?;

let prompt = context
.message
.parts
Expand All @@ -47,8 +207,9 @@ impl AgentExecutor for KitAgent {
emit.status(TaskState::Failed).await?;
return Ok(());
}
emit.status(TaskState::Working).await?;
match self
.0
.runtime
.run_cancelled(prompt, 0, Some(context.cancellation_token.clone()))
.await
{
Expand Down Expand Up @@ -109,7 +270,8 @@ pub(crate) fn dispatcher(
runtime: Arc<Runtime>,
bound: std::net::SocketAddr,
authenticated: bool,
) -> Result<JsonRpcDispatcher, Box<dyn std::error::Error>> {
sessions: Option<crate::protocols::acp::SessionRegistry>,
) -> Result<Dispatcher, Box<dyn std::error::Error>> {
let url = format!("http://{bound}");
let (security_schemes, security_requirements) = if authenticated {
let name = "bearer".to_string();
Expand All @@ -132,7 +294,11 @@ pub(crate) fn dispatcher(
let card = AgentCard {
url: None,
name: "Kit".into(),
description: "Coding agent runtime".into(),
description: if sessions.is_some() {
"Admit text into the bound interactive ACP v2 session. SendMessage returns an admission acknowledgment, not task completion or model output.".into()
} else {
"Coding agent runtime".into()
},
version: env!("CARGO_PKG_VERSION").into(),
supported_interfaces: vec![AgentInterface {
url,
Expand Down Expand Up @@ -161,9 +327,96 @@ pub(crate) fn dispatcher(
signatures: None,
};
let handler = Arc::new(
RequestHandlerBuilder::new(KitAgent(runtime))
RequestHandlerBuilder::new(KitAgent { runtime })
.with_agent_card(card)
.build()?,
);
Ok(JsonRpcDispatcher::new(handler))
Ok(Dispatcher {
inner: JsonRpcDispatcher::new(handler),
sessions,
})
}

#[cfg(test)]
#[allow(
clippy::unwrap_used,
clippy::expect_used,
clippy::panic,
clippy::unreachable,
clippy::disallowed_methods,
clippy::disallowed_macros
)]
mod tests {
use super::*;
use serde_json::json;

fn send_request() -> serde_json::Value {
json!({"jsonrpc": "2.0", "id": 1, "method": "SendMessage", "params": {
"message": {"messageId": "inbound", "role": "ROLE_USER", "parts": [{"text": "request"}]}
}})
}

#[tokio::test]
async fn a2a_direct_acknowledgment_is_a_message_not_a_completed_task() {
let registry = super::super::acp::a2a_test_support::bound_registry(Arc::new(|_prompt| {
Box::pin(async { Ok("Accepted into the bound ACP session; not completion".into()) })
}));
let acknowledgment = admit_message(&registry, send_request()).await.unwrap();
assert_eq!(acknowledgment.role, MessageRole::Agent);
assert_eq!(acknowledgment.task_id, None);
assert!(acknowledgment.text().unwrap().contains("not completion"));
assert_eq!(
acknowledgment.metadata.unwrap()["kind"],
"admission_acknowledgment"
);
}

#[tokio::test]
async fn a2a_admission_rejects_task_modes_and_missing_session_without_execution() {
let registry = crate::protocols::acp::SessionRegistry::new();
let mut request = send_request();
request["params"]["configuration"] = json!({"returnImmediately": true});
assert_eq!(
admit_message(&registry, request).await.unwrap_err().0,
-32602
);
let mut request = send_request();
request["method"] = json!("SendStreamingMessage");
assert_eq!(
admit_message(&registry, request).await.unwrap_err().0,
-32601
);
let mut request = send_request();
request["params"]["message"]["taskId"] = json!("existing-task");
assert_eq!(
admit_message(&registry, request).await.unwrap_err().0,
-32602
);
assert!(
admit_message(&registry, send_request())
.await
.unwrap_err()
.1
.contains("no bound")
);
let mut mixed = send_request();
mixed["params"]["message"]["parts"]
.as_array_mut()
.unwrap()
.push(json!({"data": {"ignored": true}}));
assert_eq!(admit_message(&registry, mixed).await.unwrap_err().0, -32602);
let mut malformed = send_request();
malformed["method"] = json!(42);
assert_eq!(
admit_message(&registry, malformed).await.unwrap_err().0,
-32600
);
let registry = super::super::acp::a2a_test_support::bound_registry(Arc::new(|_prompt| {
Box::pin(async { Err("ACP admission failed".into()) })
}));
assert_eq!(
admit_message(&registry, send_request()).await.unwrap_err(),
(-32603, "ACP admission failed".into())
);
}
}
Loading
Loading