long-running shard session via stream-json
This commit is contained in:
Damocles
parent
f4f11b78ad
commit
d4b8aa731b
1 changed files with 451 additions and 134 deletions
307
src/shard.rs
Normal file
307
src/shard.rs
Normal file
|
|
@ -0,0 +1,307 @@
|
|||
//! Long-running claude process for the shard. Spawned once, fed turns via
|
||||
//! stdin (stream-json), emits events via stdout. Survives across many turns
|
||||
//! within configured limits, then respawns fresh.
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::Stdio;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use anyhow::{Context, bail};
|
||||
use matrix_sdk::ruma::OwnedRoomId;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
|
||||
use tokio::process::{Child, ChildStdin, Command};
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::paths;
|
||||
|
||||
/// One stream-json event from claude's stdout. Only fields we care about.
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum StreamEvent {
|
||||
#[serde(rename = "system")]
|
||||
System {
|
||||
#[serde(default)]
|
||||
subtype: String,
|
||||
#[serde(default)]
|
||||
session_id: Option<String>,
|
||||
},
|
||||
#[serde(rename = "assistant")]
|
||||
Assistant {
|
||||
#[serde(default)]
|
||||
message: serde_json::Value,
|
||||
},
|
||||
#[serde(rename = "result")]
|
||||
Result {
|
||||
#[serde(default)]
|
||||
subtype: String,
|
||||
#[serde(default)]
|
||||
is_error: bool,
|
||||
#[serde(default)]
|
||||
result: String,
|
||||
#[serde(default)]
|
||||
stop_reason: Option<String>,
|
||||
},
|
||||
#[serde(rename = "rate_limit_event")]
|
||||
RateLimit {},
|
||||
#[serde(other)]
|
||||
Other,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct UserMessage<'a> {
|
||||
#[serde(rename = "type")]
|
||||
kind: &'static str,
|
||||
message: UserMessageBody<'a>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct UserMessageBody<'a> {
|
||||
role: &'static str,
|
||||
content: Vec<UserContent<'a>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct UserContent<'a> {
|
||||
#[serde(rename = "type")]
|
||||
kind: &'static str,
|
||||
text: &'a str,
|
||||
}
|
||||
|
||||
/// A live claude shard process. Owns stdin and an mpsc receiver for stdout
|
||||
/// events. When it exits or is dropped, the child is cleaned up.
|
||||
pub struct ShardSession {
|
||||
child: Child,
|
||||
stdin: ChildStdin,
|
||||
events: mpsc::Receiver<StreamEvent>,
|
||||
/// When the session was spawned.
|
||||
started: Instant,
|
||||
/// Last time a turn finished.
|
||||
last_used: Instant,
|
||||
/// Number of turns processed.
|
||||
turn_count: u32,
|
||||
/// Rooms we've sent at least one turn for in this session. Used to decide
|
||||
/// whether to include `previously_seen` context in a turn.
|
||||
pub rooms_seen: HashSet<OwnedRoomId>,
|
||||
/// Mtimes of identity/CHANGELOG files at session start - for refresh.
|
||||
mtime_snapshot: Vec<(PathBuf, std::time::SystemTime)>,
|
||||
}
|
||||
|
||||
pub struct SpawnConfig<'a> {
|
||||
pub model: &'a str,
|
||||
pub mcp_config_path: &'a Path,
|
||||
pub allowed_tools: &'a str,
|
||||
}
|
||||
|
||||
impl ShardSession {
|
||||
pub async fn spawn(cfg: SpawnConfig<'_>) -> anyhow::Result<Self> {
|
||||
let identity_dir = paths::identity_dir();
|
||||
|
||||
let mut cmd = Command::new("claude");
|
||||
cmd.args([
|
||||
"--print",
|
||||
"--input-format",
|
||||
"stream-json",
|
||||
"--output-format",
|
||||
"stream-json",
|
||||
"--model",
|
||||
cfg.model,
|
||||
"--add-dir",
|
||||
&identity_dir.to_string_lossy(),
|
||||
"--allowedTools",
|
||||
cfg.allowed_tools,
|
||||
"--mcp-config",
|
||||
&cfg.mcp_config_path.to_string_lossy(),
|
||||
]);
|
||||
cmd.current_dir(&identity_dir);
|
||||
cmd.stdin(Stdio::piped());
|
||||
cmd.stdout(Stdio::piped());
|
||||
cmd.stderr(Stdio::piped());
|
||||
cmd.kill_on_drop(true);
|
||||
|
||||
let mut child = cmd.spawn().context("spawn claude")?;
|
||||
let stdin = child.stdin.take().context("claude stdin missing")?;
|
||||
let stdout = child.stdout.take().context("claude stdout missing")?;
|
||||
let stderr = child.stderr.take().context("claude stderr missing")?;
|
||||
|
||||
let (tx, rx) = mpsc::channel(256);
|
||||
|
||||
// stdout reader: parse stream-json line by line, push events
|
||||
let tx_out = tx.clone();
|
||||
tokio::spawn(async move {
|
||||
let mut lines = BufReader::new(stdout).lines();
|
||||
while let Ok(Some(line)) = lines.next_line().await {
|
||||
if line.trim().is_empty() {
|
||||
continue;
|
||||
}
|
||||
let ev: StreamEvent = match serde_json::from_str(&line) {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
tracing::debug!(?e, line = %line, "shard: failed to parse stream-json line");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
if tx_out.send(ev).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
tracing::info!("shard: stdout stream closed");
|
||||
});
|
||||
|
||||
// stderr drainer: log claude's stderr at warn level
|
||||
tokio::spawn(async move {
|
||||
let mut lines = BufReader::new(stderr).lines();
|
||||
while let Ok(Some(line)) = lines.next_line().await {
|
||||
tracing::warn!("claude stderr: {line}");
|
||||
}
|
||||
});
|
||||
|
||||
let mtime_snapshot = snapshot_identity_mtimes();
|
||||
tracing::info!("shard: spawned");
|
||||
|
||||
Ok(Self {
|
||||
child,
|
||||
stdin,
|
||||
events: rx,
|
||||
started: Instant::now(),
|
||||
last_used: Instant::now(),
|
||||
turn_count: 0,
|
||||
rooms_seen: HashSet::new(),
|
||||
mtime_snapshot,
|
||||
})
|
||||
}
|
||||
|
||||
/// Send one user-turn JSON to claude and wait for the next `result`.
|
||||
/// Returns the result text (claude's final assistant output).
|
||||
pub async fn run_turn(&mut self, turn_body: &str) -> anyhow::Result<String> {
|
||||
let msg = UserMessage {
|
||||
kind: "user",
|
||||
message: UserMessageBody {
|
||||
role: "user",
|
||||
content: vec![UserContent {
|
||||
kind: "text",
|
||||
text: turn_body,
|
||||
}],
|
||||
},
|
||||
};
|
||||
let mut json = serde_json::to_string(&msg)?;
|
||||
json.push('\n');
|
||||
self.stdin.write_all(json.as_bytes()).await?;
|
||||
self.stdin.flush().await?;
|
||||
|
||||
// Drain events until we hit a `result` (turn end).
|
||||
loop {
|
||||
let ev = self
|
||||
.events
|
||||
.recv()
|
||||
.await
|
||||
.context("shard: stdout closed before result")?;
|
||||
match ev {
|
||||
StreamEvent::System { .. } | StreamEvent::RateLimit {} | StreamEvent::Other => {}
|
||||
StreamEvent::Assistant { message } => {
|
||||
log_assistant_text(&message);
|
||||
}
|
||||
StreamEvent::Result {
|
||||
is_error,
|
||||
result,
|
||||
stop_reason,
|
||||
..
|
||||
} => {
|
||||
self.last_used = Instant::now();
|
||||
self.turn_count += 1;
|
||||
if is_error {
|
||||
bail!("turn ended with is_error=true: {result}");
|
||||
}
|
||||
tracing::info!(
|
||||
turn = self.turn_count,
|
||||
stop_reason = ?stop_reason,
|
||||
"shard: turn complete"
|
||||
);
|
||||
return Ok(result);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn should_refresh(&mut self, idle: Duration, max_turns: u32) -> Option<&'static str> {
|
||||
if self.last_used.elapsed() > idle {
|
||||
return Some("idle gap exceeded");
|
||||
}
|
||||
if self.turn_count >= max_turns {
|
||||
return Some("max turn count reached");
|
||||
}
|
||||
if mtimes_changed(&self.mtime_snapshot) {
|
||||
return Some("identity files changed");
|
||||
}
|
||||
if self.child.try_wait().ok().flatten().is_some() {
|
||||
return Some("child process exited");
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
pub async fn shutdown(mut self) {
|
||||
// Close stdin → claude exits gracefully on EOF
|
||||
drop(self.stdin);
|
||||
// Bounded wait for clean exit
|
||||
let _ = tokio::time::timeout(Duration::from_secs(3), self.child.wait()).await;
|
||||
let _ = self.child.kill().await;
|
||||
tracing::info!(
|
||||
turns = self.turn_count,
|
||||
uptime = ?self.started.elapsed(),
|
||||
"shard: session ended"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn log_assistant_text(message: &serde_json::Value) {
|
||||
let Some(content) = message.get("content").and_then(|c| c.as_array()) else {
|
||||
return;
|
||||
};
|
||||
for item in content {
|
||||
let Some(kind) = item.get("type").and_then(|t| t.as_str()) else {
|
||||
continue;
|
||||
};
|
||||
match kind {
|
||||
"text" => {
|
||||
if let Some(text) = item.get("text").and_then(|t| t.as_str()) {
|
||||
let preview: String = text.chars().take(200).collect();
|
||||
tracing::info!("shard text: {preview}");
|
||||
tracing::trace!("shard full text: {text}");
|
||||
}
|
||||
}
|
||||
"thinking" => {
|
||||
if let Some(text) = item.get("thinking").and_then(|t| t.as_str()) {
|
||||
tracing::trace!("shard thinking: {text}");
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn snapshot_identity_mtimes() -> Vec<(PathBuf, std::time::SystemTime)> {
|
||||
let id = paths::identity_dir();
|
||||
let state = paths::state_dir();
|
||||
let candidates = [
|
||||
id.join("CLAUDE.md"),
|
||||
id.join("SYSTEM.md"),
|
||||
id.join("notes.md"),
|
||||
state.join("CHANGELOG.md"),
|
||||
];
|
||||
candidates
|
||||
.into_iter()
|
||||
.filter_map(|p| std::fs::metadata(&p).and_then(|m| m.modified()).ok().map(|t| (p, t)))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn mtimes_changed(snapshot: &[(PathBuf, std::time::SystemTime)]) -> bool {
|
||||
for (path, prev) in snapshot {
|
||||
if let Ok(now) = std::fs::metadata(path).and_then(|m| m.modified()) {
|
||||
if &now != prev {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
Loading…
Reference in a new issue