story-kit: merge 296_story_track_per_agent_token_usage_for_cost_visibility_and_optimisation

server/src/http/mcp.rs

@@ -962,6 +962,19 @@ fn handle_tools_list(id: Option<Value>) -> JsonRpcResponse {
                         },
                         "required": ["tool_name", "input"]
                     }
+                },
+                {
+                    "name": "get_token_usage",
+                    "description": "Return per-agent token usage records from the persistent log. Shows input tokens, output tokens, cache tokens, and cost in USD for each agent session. Optionally filter by story_id.",
+                    "inputSchema": {
+                        "type": "object",
+                        "properties": {
+                            "story_id": {
+                                "type": "string",
+                                "description": "Optional: filter records to a specific story (e.g. '42_my_story')"
+                            }
+                        }
+                    }
                 }
             ]
         }),
@@ -1036,6 +1049,8 @@ async fn handle_tools_call(
         "rebuild_and_restart" => tool_rebuild_and_restart(ctx).await,
         // Permission bridge (Claude Code → frontend dialog)
         "prompt_permission" => tool_prompt_permission(&args, ctx).await,
+        // Token usage
+        "get_token_usage" => tool_get_token_usage(&args, ctx),
         _ => Err(format!("Unknown tool: {tool_name}")),
     };
 
@@ -2483,6 +2498,51 @@ async fn tool_prompt_permission(args: &Value, ctx: &AppContext) -> Result<String
     }
 }
 
+fn tool_get_token_usage(args: &Value, ctx: &AppContext) -> Result<String, String> {
+    let root = ctx.state.get_project_root()?;
+    let filter_story = args.get("story_id").and_then(|v| v.as_str());
+
+    let all_records = crate::agents::token_usage::read_all(&root)?;
+    let records: Vec<_> = all_records
+        .into_iter()
+        .filter(|r| filter_story.is_none_or(|s| r.story_id == s))
+        .collect();
+
+    let total_cost: f64 = records.iter().map(|r| r.usage.total_cost_usd).sum();
+    let total_input: u64 = records.iter().map(|r| r.usage.input_tokens).sum();
+    let total_output: u64 = records.iter().map(|r| r.usage.output_tokens).sum();
+    let total_cache_create: u64 = records
+        .iter()
+        .map(|r| r.usage.cache_creation_input_tokens)
+        .sum();
+    let total_cache_read: u64 = records
+        .iter()
+        .map(|r| r.usage.cache_read_input_tokens)
+        .sum();
+
+    serde_json::to_string_pretty(&json!({
+        "records": records.iter().map(|r| json!({
+            "story_id": r.story_id,
+            "agent_name": r.agent_name,
+            "timestamp": r.timestamp,
+            "input_tokens": r.usage.input_tokens,
+            "output_tokens": r.usage.output_tokens,
+            "cache_creation_input_tokens": r.usage.cache_creation_input_tokens,
+            "cache_read_input_tokens": r.usage.cache_read_input_tokens,
+            "total_cost_usd": r.usage.total_cost_usd,
+        })).collect::<Vec<_>>(),
+        "totals": {
+            "records": records.len(),
+            "input_tokens": total_input,
+            "output_tokens": total_output,
+            "cache_creation_input_tokens": total_cache_create,
+            "cache_read_input_tokens": total_cache_read,
+            "total_cost_usd": total_cost,
+        }
+    }))
+    .map_err(|e| format!("Serialization error: {e}"))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -2592,7 +2652,8 @@ mod tests {
         assert!(names.contains(&"prompt_permission"));
         assert!(names.contains(&"get_pipeline_status"));
         assert!(names.contains(&"rebuild_and_restart"));
-        assert_eq!(tools.len(), 39);
+        assert!(names.contains(&"get_token_usage"));
+        assert_eq!(tools.len(), 40);
     }
 
     #[test]
@@ -3729,6 +3790,72 @@ mod tests {
         let _ = result;
     }
 
+    // ── tool_get_token_usage tests ────────────────────────────────
+
+    #[test]
+    fn tool_get_token_usage_empty_returns_zero_totals() {
+        let tmp = tempfile::tempdir().unwrap();
+        let ctx = test_ctx(tmp.path());
+        let result = tool_get_token_usage(&json!({}), &ctx).unwrap();
+        let parsed: Value = serde_json::from_str(&result).unwrap();
+        assert_eq!(parsed["records"].as_array().unwrap().len(), 0);
+        assert_eq!(parsed["totals"]["records"], 0);
+        assert_eq!(parsed["totals"]["total_cost_usd"], 0.0);
+    }
+
+    #[test]
+    fn tool_get_token_usage_returns_written_records() {
+        let tmp = tempfile::tempdir().unwrap();
+        let root = tmp.path();
+        let ctx = test_ctx(root);
+
+        let usage = crate::agents::TokenUsage {
+            input_tokens: 100,
+            output_tokens: 200,
+            cache_creation_input_tokens: 5000,
+            cache_read_input_tokens: 10000,
+            total_cost_usd: 1.57,
+        };
+        let record =
+            crate::agents::token_usage::build_record("42_story_foo", "coder-1", usage);
+        crate::agents::token_usage::append_record(root, &record).unwrap();
+
+        let result = tool_get_token_usage(&json!({}), &ctx).unwrap();
+        let parsed: Value = serde_json::from_str(&result).unwrap();
+        assert_eq!(parsed["records"].as_array().unwrap().len(), 1);
+        assert_eq!(parsed["records"][0]["story_id"], "42_story_foo");
+        assert_eq!(parsed["records"][0]["agent_name"], "coder-1");
+        assert_eq!(parsed["records"][0]["input_tokens"], 100);
+        assert_eq!(parsed["totals"]["records"], 1);
+        assert!((parsed["totals"]["total_cost_usd"].as_f64().unwrap() - 1.57).abs() < f64::EPSILON);
+    }
+
+    #[test]
+    fn tool_get_token_usage_filters_by_story_id() {
+        let tmp = tempfile::tempdir().unwrap();
+        let root = tmp.path();
+        let ctx = test_ctx(root);
+
+        let usage = crate::agents::TokenUsage {
+            input_tokens: 50,
+            output_tokens: 60,
+            cache_creation_input_tokens: 0,
+            cache_read_input_tokens: 0,
+            total_cost_usd: 0.5,
+        };
+        let r1 = crate::agents::token_usage::build_record("10_story_a", "coder-1", usage.clone());
+        let r2 = crate::agents::token_usage::build_record("20_story_b", "coder-2", usage);
+        crate::agents::token_usage::append_record(root, &r1).unwrap();
+        crate::agents::token_usage::append_record(root, &r2).unwrap();
+
+        let result =
+            tool_get_token_usage(&json!({"story_id": "10_story_a"}), &ctx).unwrap();
+        let parsed: Value = serde_json::from_str(&result).unwrap();
+        assert_eq!(parsed["records"].as_array().unwrap().len(), 1);
+        assert_eq!(parsed["records"][0]["story_id"], "10_story_a");
+        assert_eq!(parsed["totals"]["records"], 1);
+    }
+
     // ── tool_list_worktrees tests ─────────────────────────────────
 
     #[test]