langgraph-api 0.15.0.dev3__py3-none-any.whl → 0.15.0.dev4__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
langgraph_api/__init__.py CHANGED
@@ -1 +1 @@
1
- __version__ = "0.15.0.dev3"
1
+ __version__ = "0.15.0.dev4"
langgraph_api/api/a2a.py CHANGED
@@ -125,6 +125,30 @@ class A2UIContext:
125
125
  _A2UI_DISABLED_CONTEXT = A2UIContext(False, False, {}, [])
126
126
 
127
127
 
128
+ @dataclass(frozen=True)
129
+ class A2ATask:
130
+ task_id: str
131
+ context_id: str
132
+ run: dict[str, Any]
133
+ state: str
134
+ status_timestamp: str
135
+ task_ids_by_start_id: dict[str, str]
136
+
137
+
138
+ # Keep each core search rate-limit charge bounded.
139
+ _A2A_CONTEXT_SEARCH_LIMIT = 1000
140
+ _A2A_RUN_SCAN_LIMIT = 1000
141
+ _A2A_LIST_CONCURRENCY = 10
142
+ _TERMINAL_TASK_STATES = frozenset(
143
+ {
144
+ "TASK_STATE_COMPLETED",
145
+ "TASK_STATE_FAILED",
146
+ "TASK_STATE_CANCELED",
147
+ "TASK_STATE_REJECTED",
148
+ }
149
+ )
150
+
151
+
128
152
  def _a2ui_enabled(assistant: dict[str, Any]) -> bool:
129
153
  metadata = assistant.get("metadata")
130
154
  a2a_metadata = metadata.get("a2a") if isinstance(metadata, dict) else None
@@ -146,17 +170,30 @@ def _a2ui_context_for_run(
146
170
  )
147
171
 
148
172
 
149
- def _a2ui_run_metadata(
150
- assistant: dict[str, Any], context: A2UIContext
151
- ) -> dict[str, Any] | None:
152
- if not context.enabled:
153
- return None
173
+ def _a2a_run_metadata(
174
+ assistant: dict[str, Any],
175
+ context: A2UIContext,
176
+ *,
177
+ task_id: str | None,
178
+ task_start_id: str | None,
179
+ ) -> dict[str, Any]:
154
180
  metadata = assistant.get("metadata")
155
181
  a2a_metadata = metadata.get("a2a") if isinstance(metadata, dict) else None
182
+ inherited = (
183
+ {
184
+ key: value
185
+ for key, value in a2a_metadata.items()
186
+ if key not in {"task_id", "task_start_id"}
187
+ }
188
+ if context.enabled and isinstance(a2a_metadata, dict)
189
+ else {}
190
+ )
156
191
  return {
157
192
  "a2a": {
158
- **(a2a_metadata if isinstance(a2a_metadata, dict) else {}),
159
- "a2ui_active": context.active,
193
+ **inherited,
194
+ **({"task_id": task_id} if task_id else {}),
195
+ **({"task_start_id": task_start_id} if task_start_id else {}),
196
+ **({"a2ui_active": context.active} if context.enabled else {}),
160
197
  }
161
198
  }
162
199
 
@@ -295,6 +332,180 @@ def _parse_task_id(task_id: str) -> tuple[str, str]:
295
332
  return "", task_id
296
333
 
297
334
 
335
+ def _http_status_code(exception: Exception) -> int | None:
336
+ response = getattr(exception, "response", None)
337
+ return getattr(response, "status_code", None)
338
+
339
+
340
+ def _run_a2a_metadata(run: dict[str, Any]) -> dict[str, Any]:
341
+ metadata = run.get("metadata")
342
+ a2a = metadata.get("a2a") if isinstance(metadata, dict) else None
343
+ return a2a if isinstance(a2a, dict) else {}
344
+
345
+
346
+ def _task_id_for_run(context_id: str, run: dict[str, Any]) -> str:
347
+ stored_task_id = _run_a2a_metadata(run).get("task_id")
348
+ if isinstance(stored_task_id, str) and stored_task_id:
349
+ return stored_task_id
350
+ return _make_task_id(context_id, str(run["run_id"]))
351
+
352
+
353
+ def _task_ids_by_start_id(
354
+ runs: list[dict[str, Any]], context_id: str, assistant_id: str
355
+ ) -> dict[str, str]:
356
+ task_ids: dict[str, str] = {}
357
+ for run in runs:
358
+ if str(run.get("assistant_id")) != assistant_id:
359
+ continue
360
+ task_start_id = _run_a2a_metadata(run).get("task_start_id")
361
+ if isinstance(task_start_id, str) and task_start_id:
362
+ task_ids[task_start_id] = _task_id_for_run(context_id, run)
363
+ return task_ids
364
+
365
+
366
+ def _run_sort_key(run: dict[str, Any]) -> tuple[str, str]:
367
+ created_at = run.get("created_at")
368
+ if isinstance(created_at, datetime):
369
+ created_at = created_at.isoformat()
370
+ return str(created_at or ""), str(run.get("run_id") or "")
371
+
372
+
373
+ def _latest_run_id(runs: list[dict[str, Any]]) -> str | None:
374
+ return str(max(runs, key=_run_sort_key).get("run_id")) if runs else None
375
+
376
+
377
+ def _runs_for_assistant(
378
+ runs: list[dict[str, Any]], assistant_id: str
379
+ ) -> list[dict[str, Any]]:
380
+ return sorted(
381
+ (run for run in runs if str(run.get("assistant_id")) == assistant_id),
382
+ key=_run_sort_key,
383
+ reverse=True,
384
+ )
385
+
386
+
387
+ async def _list_context_runs(
388
+ client: LangGraphClient,
389
+ *,
390
+ context_id: str,
391
+ headers: Headers | dict[str, Any] | None,
392
+ ) -> list[dict[str, Any]]:
393
+ runs = await client.runs.list(
394
+ context_id,
395
+ limit=_A2A_RUN_SCAN_LIMIT,
396
+ headers=headers,
397
+ )
398
+ if len(runs) == _A2A_RUN_SCAN_LIMIT:
399
+ await logger.awarning(
400
+ "A2A context run scan reached its limit; task state or history may be truncated",
401
+ context_id=context_id,
402
+ run_scan_limit=_A2A_RUN_SCAN_LIMIT,
403
+ )
404
+ return runs
405
+
406
+
407
+ def _lg_status_to_a2a_state(lg_status: str) -> str:
408
+ """Map a LangGraph run status to an A2A task state."""
409
+ mapping = {
410
+ "pending": "TASK_STATE_SUBMITTED",
411
+ "running": "TASK_STATE_WORKING",
412
+ "success": "TASK_STATE_COMPLETED",
413
+ # LangGraph uses this for an externally canceled run. Graph-level HITL
414
+ # interrupts finish the run successfully and interrupt the thread.
415
+ "interrupted": "TASK_STATE_CANCELED",
416
+ "error": "TASK_STATE_FAILED",
417
+ "timeout": "TASK_STATE_FAILED",
418
+ }
419
+ return mapping.get(lg_status, "TASK_STATE_SUBMITTED")
420
+
421
+
422
+ def _task_state_for_run(
423
+ run: dict[str, Any], thread: dict[str, Any] | None = None
424
+ ) -> str:
425
+ state = _lg_status_to_a2a_state(run.get("status", "unknown"))
426
+ if (
427
+ run.get("status") == "success"
428
+ and thread is not None
429
+ and thread.get("status") == "interrupted"
430
+ ):
431
+ return "TASK_STATE_INPUT_REQUIRED"
432
+ return state
433
+
434
+
435
+ def _run_timestamp(run: dict[str, Any]) -> str:
436
+ timestamp = run.get("updated_at") or run.get("created_at")
437
+ if isinstance(timestamp, datetime):
438
+ return timestamp.isoformat()
439
+ if isinstance(timestamp, str):
440
+ return timestamp
441
+ return datetime.now(UTC).isoformat()
442
+
443
+
444
+ async def _resolve_task(
445
+ client: LangGraphClient,
446
+ *,
447
+ task_id: str,
448
+ context_id: str | None,
449
+ assistant_id: str,
450
+ headers: Headers | dict[str, Any] | None,
451
+ ) -> A2ATask | None:
452
+ parsed_context_id, anchor_run_id = _parse_task_id(task_id)
453
+ if parsed_context_id and context_id and parsed_context_id != context_id:
454
+ return None
455
+ resolved_context_id = context_id or parsed_context_id
456
+ if not resolved_context_id:
457
+ return None
458
+
459
+ try:
460
+ anchor_run, thread = await asyncio.gather(
461
+ client.runs.get(
462
+ thread_id=resolved_context_id,
463
+ run_id=anchor_run_id,
464
+ headers=headers,
465
+ ),
466
+ client.threads.get(resolved_context_id, headers=headers),
467
+ )
468
+ except Exception as exc:
469
+ if _http_status_code(exc) in {400, 404, 422}:
470
+ return None
471
+ raise
472
+ if str(anchor_run.get("assistant_id")) != assistant_id:
473
+ return None
474
+
475
+ anchor_stored_task_id = _run_a2a_metadata(anchor_run).get("task_id")
476
+ if (
477
+ isinstance(anchor_stored_task_id, str)
478
+ and anchor_stored_task_id
479
+ and anchor_stored_task_id != task_id
480
+ ):
481
+ return None
482
+
483
+ runs = await _list_context_runs(
484
+ client, context_id=resolved_context_id, headers=headers
485
+ )
486
+ latest_run = anchor_run
487
+ context_latest_run_id = _latest_run_id(runs)
488
+ is_context_latest = False
489
+ for run in _runs_for_assistant(runs, assistant_id):
490
+ if str(run.get("run_id")) == anchor_run_id or (
491
+ _run_a2a_metadata(run).get("task_id") == task_id
492
+ ):
493
+ latest_run = run
494
+ is_context_latest = str(run.get("run_id")) == context_latest_run_id
495
+ break
496
+
497
+ return A2ATask(
498
+ task_id=task_id,
499
+ context_id=resolved_context_id,
500
+ run=latest_run,
501
+ state=_task_state_for_run(latest_run, thread if is_context_latest else None),
502
+ status_timestamp=_run_timestamp(latest_run),
503
+ task_ids_by_start_id=_task_ids_by_start_id(
504
+ runs, resolved_context_id, assistant_id
505
+ ),
506
+ )
507
+
508
+
298
509
  async def _get_assistant(
299
510
  assistant_id: str, headers: Headers | dict[str, Any] | None
300
511
  ) -> dict[str, Any]:
@@ -532,17 +743,10 @@ async def _maybe_promote_resume_to_command(
532
743
 
533
744
  if not context_id:
534
745
  return None, None, input_content
746
+ if not task_id:
747
+ return None, None, input_content
535
748
  if not await _is_thread_interrupted(client, context_id, headers):
536
749
  return None, None, input_content
537
- if not task_id:
538
- return (
539
- None,
540
- {
541
- "code": ERROR_CODE_INVALID_PARAMS,
542
- "message": "contextId and taskId are required when resuming a task",
543
- },
544
- input_content,
545
- )
546
750
 
547
751
  if resume_source == "text" and "messages" in input_content:
548
752
  input_content = dict(input_content)
@@ -593,6 +797,8 @@ def _process_a2a_message_parts(
593
797
  message_id: str,
594
798
  *,
595
799
  a2ui: A2UIContext,
800
+ task_id: str | None = None,
801
+ task_start_id: str | None = None,
596
802
  ) -> dict[str, Any]:
597
803
  """Convert A2A message parts to LangChain messages format.
598
804
 
@@ -699,8 +905,10 @@ def _process_a2a_message_parts(
699
905
  }
700
906
  )
701
907
 
702
- if a2ui.metadata or a2ui.extensions:
908
+ if task_id or task_start_id or a2ui.metadata or a2ui.extensions:
703
909
  a2a_data = {
910
+ **({"task_id": task_id} if task_id else {}),
911
+ **({"task_start_id": task_start_id} if task_start_id else {}),
704
912
  **({"metadata": a2ui.metadata} if a2ui.metadata else {}),
705
913
  **({"extensions": a2ui.extensions} if a2ui.extensions else {}),
706
914
  }
@@ -1241,62 +1449,6 @@ def _map_runs_create_error_to_rpc(
1241
1449
  }
1242
1450
 
1243
1451
 
1244
- def _map_runs_get_error_to_rpc(
1245
- exception: Exception, task_id: str, thread_id: str
1246
- ) -> dict[str, Any]:
1247
- """Map runs.get() exceptions to A2A JSON-RPC error responses.
1248
-
1249
- Args:
1250
- exception: Exception from runs.get()
1251
- task_id: The task/run ID that was requested
1252
- thread_id: The thread ID that was requested
1253
-
1254
- Returns:
1255
- A2A error response dictionary
1256
- """
1257
- if hasattr(exception, "response") and hasattr(exception.response, "status_code"):
1258
- status_code = exception.response.status_code
1259
- error_text = str(exception)
1260
-
1261
- status_code_handlers = {
1262
- 404: {
1263
- "error": {
1264
- "code": ERROR_CODE_TASK_NOT_FOUND,
1265
- "message": f"Task '{task_id}' not found in thread '{thread_id}'",
1266
- }
1267
- },
1268
- 400: {
1269
- "error": {
1270
- "code": ERROR_CODE_INVALID_PARAMS,
1271
- "message": f"Invalid request: {error_text}",
1272
- }
1273
- },
1274
- 403: {
1275
- "error": {
1276
- "code": ERROR_CODE_INVALID_PARAMS,
1277
- "message": "Access denied to task",
1278
- }
1279
- },
1280
- }
1281
-
1282
- return status_code_handlers.get(
1283
- status_code,
1284
- {
1285
- "error": {
1286
- "code": ERROR_CODE_INVALID_PARAMS,
1287
- "message": f"Failed to get task: {error_text}",
1288
- }
1289
- },
1290
- )
1291
-
1292
- return {
1293
- "error": {
1294
- "code": ERROR_CODE_INTERNAL_ERROR,
1295
- "message": "Internal server error",
1296
- }
1297
- }
1298
-
1299
-
1300
1452
  def _convert_messages_to_a2a_format(
1301
1453
  messages: list[dict[str, Any]],
1302
1454
  task_id: str,
@@ -1304,6 +1456,7 @@ def _convert_messages_to_a2a_format(
1304
1456
  *,
1305
1457
  history_length: int | None = None,
1306
1458
  a2ui: A2UIContext,
1459
+ task_ids_by_start_id: dict[str, str] | None = None,
1307
1460
  ) -> list[dict[str, Any]]:
1308
1461
  """Convert LangChain messages to A2A message format.
1309
1462
 
@@ -1315,7 +1468,7 @@ def _convert_messages_to_a2a_format(
1315
1468
 
1316
1469
  Args:
1317
1470
  messages: List of LangChain messages
1318
- task_id: The task ID to assign to all messages
1471
+ task_id: The fallback task ID for legacy messages without stored provenance
1319
1472
  context_id: The context ID to assign to all messages
1320
1473
  history_length: Maximum number of visible A2A messages to return.
1321
1474
 
@@ -1325,10 +1478,40 @@ def _convert_messages_to_a2a_format(
1325
1478
 
1326
1479
  history_a2ui = replace(a2ui, drop_invalid=True)
1327
1480
 
1481
+ # A task-stamped user message starts a new provenance segment. Agent and
1482
+ # tool messages that follow belong to that task until the next stamped turn.
1483
+ current_task_id: str | None = task_id
1484
+ known_task_ids = (
1485
+ set(task_ids_by_start_id.values()) if task_ids_by_start_id is not None else None
1486
+ )
1487
+
1328
1488
  # Convert each LangChain message to A2A format
1329
1489
  a2a_messages = []
1330
1490
  for msg in messages:
1331
1491
  if isinstance(msg, dict):
1492
+ additional_kwargs = msg.get("additional_kwargs")
1493
+ stored_a2a = (
1494
+ additional_kwargs.get("a2a")
1495
+ if isinstance(additional_kwargs, dict)
1496
+ else None
1497
+ )
1498
+ stored_task_id = (
1499
+ stored_a2a.get("task_id") if isinstance(stored_a2a, dict) else None
1500
+ )
1501
+ if isinstance(stored_task_id, str) and stored_task_id:
1502
+ current_task_id = (
1503
+ stored_task_id
1504
+ if known_task_ids is None or stored_task_id in known_task_ids
1505
+ else None
1506
+ )
1507
+ elif isinstance(stored_a2a, dict) and task_ids_by_start_id is not None:
1508
+ task_start_id = stored_a2a.get("task_start_id")
1509
+ if isinstance(task_start_id, str):
1510
+ current_task_id = task_ids_by_start_id.get(task_start_id)
1511
+
1512
+ if current_task_id is None:
1513
+ continue
1514
+
1332
1515
  a2a_role = _a2a_role_for_message(msg)
1333
1516
 
1334
1517
  if a2a_role is None:
@@ -1355,7 +1538,7 @@ def _convert_messages_to_a2a_format(
1355
1538
  "role": a2a_role,
1356
1539
  "parts": parts,
1357
1540
  "messageId": msg.get("id") or str(uuid7()),
1358
- "taskId": task_id,
1541
+ "taskId": current_task_id,
1359
1542
  "contextId": context_id,
1360
1543
  }
1361
1544
  a2a_messages.append(a2a_message)
@@ -1374,6 +1557,7 @@ async def _create_task_response(
1374
1557
  assistant_id: str,
1375
1558
  a2ui: A2UIContext,
1376
1559
  history_length: int | None = None,
1560
+ task_ids_by_start_id: dict[str, str] | None = None,
1377
1561
  ) -> dict[str, Any]:
1378
1562
  """Create A2A Task response structure for both success and failure cases.
1379
1563
 
@@ -1395,6 +1579,7 @@ async def _create_task_response(
1395
1579
  context_id,
1396
1580
  history_length=history_length,
1397
1581
  a2ui=a2ui,
1582
+ task_ids_by_start_id=task_ids_by_start_id,
1398
1583
  )
1399
1584
 
1400
1585
  base_task: dict[str, Any] = {
@@ -1685,6 +1870,14 @@ async def handle_jsonrpc_request(
1685
1870
  Returns:
1686
1871
  JSON-RPC response
1687
1872
  """
1873
+ try:
1874
+ assistant_id = str(uuid.UUID(assistant_id))
1875
+ except ValueError:
1876
+ return create_jsonrpc_error_response(
1877
+ ERROR_CODE_INVALID_PARAMS,
1878
+ "Invalid assistant ID: must be a UUID",
1879
+ message["id"],
1880
+ )
1688
1881
  method = message["method"]
1689
1882
  params = message.get("params", {})
1690
1883
  # Route to appropriate A2A method handler
@@ -1695,9 +1888,9 @@ async def handle_jsonrpc_request(
1695
1888
  elif method == "GetTask":
1696
1889
  result_or_error = await handle_tasks_get(request, params, assistant_id)
1697
1890
  elif method == "CancelTask":
1698
- result_or_error = await handle_tasks_cancel(request, params)
1891
+ result_or_error = await handle_tasks_cancel(request, params, assistant_id)
1699
1892
  elif method == "ListTasks":
1700
- result_or_error = await handle_list_tasks(request, params)
1893
+ result_or_error = await handle_list_tasks(request, params, assistant_id)
1701
1894
  elif method == "GetExtendedAgentCard":
1702
1895
  result_or_error = await handle_get_extended_card(request, assistant_id)
1703
1896
  else:
@@ -1837,6 +2030,57 @@ async def handle_message_send(
1837
2030
  }
1838
2031
  }
1839
2032
 
2033
+ requested_context_id = message.get("contextId")
2034
+ existing_task_id = message.get("taskId")
2035
+ if requested_context_id is not None and not isinstance(
2036
+ requested_context_id, str
2037
+ ):
2038
+ return {
2039
+ "error": {
2040
+ "code": ERROR_CODE_INVALID_PARAMS,
2041
+ "message": "contextId must be a string",
2042
+ }
2043
+ }
2044
+ if existing_task_id is not None and not isinstance(existing_task_id, str):
2045
+ return {
2046
+ "error": {
2047
+ "code": ERROR_CODE_INVALID_PARAMS,
2048
+ "message": "taskId must be a string",
2049
+ }
2050
+ }
2051
+
2052
+ if existing_task_id is None:
2053
+ context_id = requested_context_id or str(uuid.uuid4())
2054
+ task_id = None
2055
+ task_start_id = str(uuid7())
2056
+ task_ids_by_start_id: dict[str, str] = {}
2057
+ else:
2058
+ task = await _resolve_task(
2059
+ client,
2060
+ task_id=existing_task_id,
2061
+ context_id=requested_context_id,
2062
+ assistant_id=assistant_id,
2063
+ headers=request.headers,
2064
+ )
2065
+ if task is None:
2066
+ return {
2067
+ "error": {
2068
+ "code": ERROR_CODE_TASK_NOT_FOUND,
2069
+ "message": f"Task not found: {existing_task_id}",
2070
+ }
2071
+ }
2072
+ if task.state in _TERMINAL_TASK_STATES:
2073
+ return {
2074
+ "error": {
2075
+ "code": ERROR_CODE_UNSUPPORTED_OPERATION,
2076
+ "message": "A terminal task cannot accept another message",
2077
+ }
2078
+ }
2079
+ context_id = task.context_id
2080
+ task_id = task.task_id
2081
+ task_start_id = None
2082
+ task_ids_by_start_id = task.task_ids_by_start_id
2083
+
1840
2084
  # Process A2A message parts into LangChain messages format
1841
2085
  try:
1842
2086
  a2ui = _a2ui_message_context(message, assistant)
@@ -1846,6 +2090,8 @@ async def handle_message_send(
1846
2090
  message_role,
1847
2091
  message_id,
1848
2092
  a2ui=a2ui,
2093
+ task_id=task_id,
2094
+ task_start_id=task_start_id,
1849
2095
  )
1850
2096
  except A2UIValidationError as e:
1851
2097
  return {
@@ -1862,12 +2108,10 @@ async def handle_message_send(
1862
2108
  }
1863
2109
  }
1864
2110
 
1865
- context_id = message.get("contextId")
1866
- # Check if this is a continuation (taskId provided in message)
1867
- existing_task_id = message.get("taskId")
1868
-
1869
2111
  # Extract and validate command (LangGraph extension for resuming interrupts)
1870
- command, command_error = _extract_and_validate_command(message, context_id)
2112
+ command, command_error = _extract_and_validate_command(
2113
+ message, requested_context_id
2114
+ )
1871
2115
  if command_error:
1872
2116
  return {"error": command_error}
1873
2117
  if command is not None and command.get("resume") and existing_task_id is None:
@@ -1885,7 +2129,7 @@ async def handle_message_send(
1885
2129
  ) = await _maybe_promote_resume_to_command(
1886
2130
  client=client,
1887
2131
  parts=parts,
1888
- context_id=context_id,
2132
+ context_id=context_id if existing_task_id else requested_context_id,
1889
2133
  task_id=existing_task_id,
1890
2134
  input_content=input_content,
1891
2135
  headers=request.headers,
@@ -1899,10 +2143,6 @@ async def handle_message_send(
1899
2143
  task_id=existing_task_id,
1900
2144
  )
1901
2145
 
1902
- # If no contextId provided, generate a UUID so we don't pass None to runs.create
1903
- if context_id is None:
1904
- context_id = str(uuid.uuid4())
1905
-
1906
2146
  try:
1907
2147
  run = await client.runs.create(
1908
2148
  thread_id=context_id,
@@ -1910,7 +2150,12 @@ async def handle_message_send(
1910
2150
  input=input_content,
1911
2151
  command=command,
1912
2152
  context=run_context,
1913
- metadata=_a2ui_run_metadata(assistant, a2ui),
2153
+ metadata=_a2a_run_metadata(
2154
+ assistant,
2155
+ a2ui,
2156
+ task_id=task_id,
2157
+ task_start_id=task_start_id,
2158
+ ),
1914
2159
  if_not_exists="create",
1915
2160
  headers=request.headers,
1916
2161
  )
@@ -1920,17 +2165,31 @@ async def handle_message_send(
1920
2165
  raise
1921
2166
  return error_response
1922
2167
 
2168
+ context_id = str(run["thread_id"])
2169
+ if task_id is None:
2170
+ task_id = _make_task_id(context_id, str(run["run_id"]))
2171
+ try:
2172
+ runs = await _list_context_runs(
2173
+ client, context_id=context_id, headers=request.headers
2174
+ )
2175
+ task_ids_by_start_id = _task_ids_by_start_id(
2176
+ runs, context_id, assistant_id
2177
+ )
2178
+ except Exception:
2179
+ await logger.awarning(
2180
+ "Failed to resolve A2A history provenance",
2181
+ context_id=context_id,
2182
+ exc_info=True,
2183
+ )
2184
+ if task_start_id:
2185
+ task_ids_by_start_id.setdefault(task_start_id, task_id)
2186
+
1923
2187
  result = await client.runs.join(
1924
2188
  thread_id=run["thread_id"],
1925
2189
  run_id=run["run_id"],
1926
2190
  headers=request.headers,
1927
2191
  )
1928
2192
 
1929
- context_id = str(run["thread_id"])
1930
- # If continuing an existing task, preserve the original task_id
1931
- # Otherwise create a new composite task_id
1932
- task_id = existing_task_id or _make_task_id(context_id, run["run_id"])
1933
-
1934
2193
  return await _create_task_response(
1935
2194
  task_id=task_id,
1936
2195
  context_id=context_id,
@@ -1938,6 +2197,7 @@ async def handle_message_send(
1938
2197
  assistant_id=assistant_id,
1939
2198
  history_length=history_length,
1940
2199
  a2ui=a2ui,
2200
+ task_ids_by_start_id=task_ids_by_start_id,
1941
2201
  )
1942
2202
 
1943
2203
  except InvalidAgentResponseError as e:
@@ -2012,24 +2272,36 @@ async def handle_tasks_get(
2012
2272
  "message": "Missing required parameter: id (task_id)",
2013
2273
  }
2014
2274
  }
2275
+ if not isinstance(task_id_raw, str) or (
2276
+ context_id_param is not None and not isinstance(context_id_param, str)
2277
+ ):
2278
+ return {
2279
+ "error": {
2280
+ "code": ERROR_CODE_INVALID_PARAMS,
2281
+ "message": "Task id and contextId must be strings",
2282
+ }
2283
+ }
2015
2284
 
2016
- # Parse composite task_id to extract context_id and run_id
2017
- parsed_context_id, run_id = _parse_task_id(task_id_raw)
2018
-
2019
- # Use contextId from params if provided, otherwise from task_id
2020
- context_id = context_id_param or parsed_context_id
2021
-
2022
- if not context_id:
2023
- # If task_id isn't a composite ID and no contextId provided, task doesn't exist
2285
+ task = await _resolve_task(
2286
+ client,
2287
+ task_id=task_id_raw,
2288
+ context_id=context_id_param,
2289
+ assistant_id=assistant_id,
2290
+ headers=request.headers,
2291
+ )
2292
+ if task is None:
2293
+ location = f" in context '{context_id_param}'" if context_id_param else ""
2024
2294
  return {
2025
2295
  "error": {
2026
2296
  "code": ERROR_CODE_TASK_NOT_FOUND,
2027
- "message": f"Task not found: {task_id_raw}",
2297
+ "message": f"Task not found: {task_id_raw}{location}",
2028
2298
  }
2029
2299
  }
2030
2300
 
2031
2301
  # Keep original task_id for A2A response (preserve what was sent/received)
2032
- task_id = task_id_raw
2302
+ task_id = task.task_id
2303
+ context_id = task.context_id
2304
+ run_info = task.run
2033
2305
 
2034
2306
  if history_error := _validate_history_length(history_length):
2035
2307
  return {
@@ -2039,49 +2311,10 @@ async def handle_tasks_get(
2039
2311
  }
2040
2312
  }
2041
2313
 
2042
- try:
2043
- # TODO: fix the N+1 query issue
2044
- assistant, run_info, thread_info = await asyncio.gather(
2045
- _get_assistant(assistant_id, request.headers),
2046
- client.runs.get(
2047
- thread_id=context_id,
2048
- run_id=run_id,
2049
- headers=request.headers,
2050
- ),
2051
- client.threads.get(
2052
- thread_id=context_id,
2053
- headers=request.headers,
2054
- ),
2055
- )
2056
- except Exception as e:
2057
- error_response = _map_runs_get_error_to_rpc(e, run_id, context_id)
2058
- if error_response.get("error", {}).get("code") == ERROR_CODE_INTERNAL_ERROR:
2059
- # For unmapped errors, re-raise to be caught by outer exception handler
2060
- raise
2061
- return error_response
2314
+ assistant = await _get_assistant(assistant_id, request.headers)
2062
2315
 
2063
2316
  lg_status = run_info.get("status", "unknown")
2064
-
2065
- if lg_status == "pending":
2066
- a2a_state = "TASK_STATE_SUBMITTED"
2067
- elif lg_status == "running":
2068
- a2a_state = "TASK_STATE_WORKING"
2069
- elif lg_status == "success":
2070
- # Hack hack: if the thread **at present** is interrupted, assume
2071
- # the run also is interrupted
2072
- if thread_info.get("status") == "interrupted":
2073
- a2a_state = "TASK_STATE_INPUT_REQUIRED"
2074
- else:
2075
- # Inspect whether there are next tasks
2076
- a2a_state = "TASK_STATE_COMPLETED"
2077
- elif (
2078
- lg_status == "interrupted"
2079
- ): # Note that this is if you interrupt FROM the outside (i.e., with double texting)
2080
- a2a_state = "TASK_STATE_INPUT_REQUIRED"
2081
- elif lg_status in ["error", "timeout"]:
2082
- a2a_state = "TASK_STATE_FAILED"
2083
- else:
2084
- a2a_state = "TASK_STATE_SUBMITTED"
2317
+ a2a_state = task.state
2085
2318
 
2086
2319
  thread_history = []
2087
2320
  if history_length != 0:
@@ -2098,6 +2331,7 @@ async def handle_tasks_get(
2098
2331
  a2ui=_a2ui_context_for_run(
2099
2332
  run_info, enabled=_a2ui_enabled(assistant)
2100
2333
  ),
2334
+ task_ids_by_start_id=task.task_ids_by_start_id,
2101
2335
  )
2102
2336
  except A2UIValidationError:
2103
2337
  raise
@@ -2120,6 +2354,7 @@ async def handle_tasks_get(
2120
2354
  "contextId": context_id,
2121
2355
  "status": {
2122
2356
  "state": a2a_state,
2357
+ "timestamp": task.status_timestamp,
2123
2358
  },
2124
2359
  }
2125
2360
  if history_length != 0:
@@ -2160,7 +2395,7 @@ async def handle_tasks_get(
2160
2395
 
2161
2396
 
2162
2397
  async def handle_tasks_cancel(
2163
- request: ApiRequest, params: dict[str, Any]
2398
+ request: ApiRequest, params: dict[str, Any], assistant_id: str
2164
2399
  ) -> dict[str, Any]:
2165
2400
  """Handle tasks/cancel requests to cancel running tasks.
2166
2401
 
@@ -2189,87 +2424,111 @@ async def handle_tasks_cancel(
2189
2424
  "message": "Missing required parameter: id (task_id)",
2190
2425
  }
2191
2426
  }
2192
-
2193
- # Parse composite task_id to extract context_id and run_id
2194
- parsed_context_id, run_id = _parse_task_id(task_id_raw)
2195
- context_id = context_id_param or parsed_context_id
2196
-
2197
- if not context_id:
2198
- # If task_id isn't a composite ID and no contextId provided, task doesn't exist
2427
+ if not isinstance(task_id_raw, str) or (
2428
+ context_id_param is not None and not isinstance(context_id_param, str)
2429
+ ):
2199
2430
  return {
2200
2431
  "error": {
2201
- "code": ERROR_CODE_TASK_NOT_FOUND,
2202
- "message": f"Task not found: {task_id_raw}",
2432
+ "code": ERROR_CODE_INVALID_PARAMS,
2433
+ "message": "Task id and contextId must be strings",
2203
2434
  }
2204
2435
  }
2205
2436
 
2206
- # Check if the task exists first
2207
2437
  try:
2208
- run_info = await client.runs.get(
2209
- thread_id=context_id,
2210
- run_id=run_id,
2438
+ task = await _resolve_task(
2439
+ client,
2440
+ task_id=task_id_raw,
2441
+ context_id=context_id_param,
2442
+ assistant_id=assistant_id,
2211
2443
  headers=request.headers,
2212
2444
  )
2213
- except Exception as e:
2214
- # Check if it's a 404 error
2215
- if (
2216
- hasattr(e, "response")
2217
- and hasattr(e.response, "status_code")
2218
- and e.response.status_code == 404
2219
- ):
2220
- return {
2221
- "error": {
2222
- "code": ERROR_CODE_TASK_NOT_FOUND,
2223
- "message": f"Task not found: {task_id_raw}",
2224
- }
2225
- }
2226
- # For other errors, return internal error
2445
+ except Exception:
2446
+ logger.exception("Failed to resolve A2A task for cancellation")
2227
2447
  return {
2228
2448
  "error": {
2229
2449
  "code": ERROR_CODE_INTERNAL_ERROR,
2230
- "message": f"Failed to check task status: {e!s}",
2450
+ "message": "Failed to check task status",
2231
2451
  }
2232
2452
  }
2233
2453
 
2234
- # Check if the task is in a cancelable state
2235
- lg_status = run_info.get("status", "unknown")
2454
+ if task is None:
2455
+ return {
2456
+ "error": {
2457
+ "code": ERROR_CODE_TASK_NOT_FOUND,
2458
+ "message": f"Task not found: {task_id_raw}",
2459
+ }
2460
+ }
2236
2461
 
2237
- # If task is already in a terminal state, return it as "canceled" per A2A spec
2238
- # The spec expects idempotent behavior - tasks/cancel always returns canceled state
2239
- if lg_status not in ("pending", "running"):
2240
- task_response = {
2241
- "kind": "task",
2242
- "id": task_id_raw,
2243
- "contextId": context_id,
2244
- "status": {
2245
- "state": "TASK_STATE_CANCELED",
2246
- "message": {
2247
- "kind": "message",
2248
- "role": "ROLE_AGENT",
2249
- "parts": [
2250
- {
2251
- "kind": "text",
2252
- "text": f"Task cancel acknowledged (was: {lg_status})",
2253
- }
2254
- ],
2255
- "messageId": str(uuid.uuid4()),
2256
- "taskId": task_id_raw,
2257
- },
2258
- },
2462
+ if task.state in _TERMINAL_TASK_STATES:
2463
+ return {
2464
+ "error": {
2465
+ "code": ERROR_CODE_TASK_NOT_CANCELABLE,
2466
+ "message": "Task is not cancelable in its current state",
2467
+ }
2259
2468
  }
2260
- return {"result": task_response}
2261
2469
 
2262
- # Cancel the run
2470
+ marker_run_id: str | None = None
2263
2471
  try:
2264
- await client.runs.cancel(
2265
- thread_id=context_id,
2266
- run_id=run_id,
2267
- wait=True, # Wait for cancellation to complete
2268
- action="interrupt",
2269
- headers=request.headers,
2270
- )
2472
+ run_id = str(task.run["run_id"])
2473
+ if task.state == "TASK_STATE_INPUT_REQUIRED":
2474
+ canonical_task_id = _task_id_for_run(task.context_id, task.run)
2475
+ marker_a2a = {**_run_a2a_metadata(task.run), "task_id": canonical_task_id}
2476
+ marker = await client.runs.create(
2477
+ thread_id=task.context_id,
2478
+ assistant_id=assistant_id,
2479
+ metadata={"a2a": marker_a2a},
2480
+ # Keep the marker pending until it is persisted as canceled below.
2481
+ after_seconds=3600,
2482
+ headers=request.headers,
2483
+ )
2484
+ marker_run_id = str(marker["run_id"])
2485
+ await client.runs.cancel(
2486
+ thread_id=task.context_id,
2487
+ run_id=marker_run_id,
2488
+ wait=True,
2489
+ action="interrupt",
2490
+ headers=request.headers,
2491
+ )
2492
+ await client.threads.update_state(
2493
+ task.context_id,
2494
+ None,
2495
+ as_node="__end__",
2496
+ headers=request.headers,
2497
+ )
2498
+ elif task.run.get("status") not in {"pending", "running"}:
2499
+ return {
2500
+ "error": {
2501
+ "code": ERROR_CODE_TASK_NOT_CANCELABLE,
2502
+ "message": "Task has no active run to cancel",
2503
+ }
2504
+ }
2505
+ else:
2506
+ await client.runs.cancel(
2507
+ thread_id=task.context_id,
2508
+ run_id=run_id,
2509
+ wait=True,
2510
+ action="interrupt",
2511
+ headers=request.headers,
2512
+ )
2271
2513
  except Exception as e:
2272
- await logger.aerror(f"Failed to cancel run {run_id}: {e!s}", exc_info=True)
2514
+ if marker_run_id is not None:
2515
+ try:
2516
+ await client.runs.delete(
2517
+ task.context_id,
2518
+ marker_run_id,
2519
+ headers=request.headers,
2520
+ )
2521
+ except Exception:
2522
+ logger.exception(
2523
+ "Failed to roll back A2A cancellation marker",
2524
+ task_id=task_id_raw,
2525
+ marker_run_id=marker_run_id,
2526
+ )
2527
+ await logger.aerror(
2528
+ "Failed to cancel A2A task",
2529
+ task_id=task_id_raw,
2530
+ exc_info=True,
2531
+ )
2273
2532
  return {
2274
2533
  "error": {
2275
2534
  "code": ERROR_CODE_INTERNAL_ERROR,
@@ -2281,9 +2540,10 @@ async def handle_tasks_cancel(
2281
2540
  task_response = {
2282
2541
  "kind": "task",
2283
2542
  "id": task_id_raw,
2284
- "contextId": context_id,
2543
+ "contextId": task.context_id,
2285
2544
  "status": {
2286
2545
  "state": "TASK_STATE_CANCELED",
2546
+ "timestamp": datetime.now(UTC).isoformat(),
2287
2547
  "message": {
2288
2548
  "kind": "message",
2289
2549
  "role": "ROLE_AGENT",
@@ -2297,21 +2557,8 @@ async def handle_tasks_cancel(
2297
2557
  return {"result": task_response}
2298
2558
 
2299
2559
 
2300
- def _lg_status_to_a2a_state(lg_status: str) -> str:
2301
- """Map a LangGraph run status to an A2A task state."""
2302
- mapping = {
2303
- "pending": "TASK_STATE_SUBMITTED",
2304
- "running": "TASK_STATE_WORKING",
2305
- "success": "TASK_STATE_COMPLETED",
2306
- "interrupted": "TASK_STATE_INPUT_REQUIRED",
2307
- "error": "TASK_STATE_FAILED",
2308
- "timeout": "TASK_STATE_FAILED",
2309
- }
2310
- return mapping.get(lg_status, "TASK_STATE_SUBMITTED")
2311
-
2312
-
2313
2560
  async def handle_list_tasks(
2314
- request: ApiRequest, params: dict[str, Any]
2561
+ request: ApiRequest, params: dict[str, Any], assistant_id: str
2315
2562
  ) -> dict[str, Any]:
2316
2563
  """Handle ListTasks requests to list tasks with filtering and pagination.
2317
2564
 
@@ -2413,43 +2660,84 @@ async def handle_list_tasks(
2413
2660
 
2414
2661
  include_artifacts = params.get("includeArtifacts", False)
2415
2662
  context_id = params.get("contextId")
2663
+ if context_id is not None and not isinstance(context_id, str):
2664
+ return {
2665
+ "error": {
2666
+ "code": ERROR_CODE_INVALID_PARAMS,
2667
+ "message": "contextId must be a string.",
2668
+ }
2669
+ }
2416
2670
 
2417
2671
  client = _client()
2418
2672
 
2419
2673
  try:
2420
2674
  # Determine which threads to search
2421
- if context_id:
2422
- thread_ids = [context_id]
2675
+ if context_id is not None:
2676
+ try:
2677
+ thread = await client.threads.get(context_id, headers=request.headers)
2678
+ except Exception as exc:
2679
+ if _http_status_code(exc) not in {400, 404, 422}:
2680
+ raise
2681
+ threads = []
2682
+ else:
2683
+ threads = [thread]
2423
2684
  else:
2424
2685
  threads = await client.threads.search(
2425
- limit=1000,
2686
+ limit=_A2A_CONTEXT_SEARCH_LIMIT,
2426
2687
  headers=request.headers,
2427
2688
  )
2428
- thread_ids = [t["thread_id"] for t in threads]
2429
2689
 
2430
- # Collect all runs from matching threads
2431
- all_tasks: list[dict[str, Any]] = []
2690
+ # ponytail: runs are thread-scoped; bounded fan-out is the smallest safe
2691
+ # option until the core API supports bulk run searches across threads.
2692
+ semaphore = asyncio.Semaphore(_A2A_LIST_CONCURRENCY)
2693
+
2694
+ async def list_thread_runs(thread: dict[str, Any]) -> list[dict[str, Any]]:
2695
+ tid = str(thread["thread_id"])
2696
+ async with semaphore:
2697
+ try:
2698
+ return await _list_context_runs(
2699
+ client, context_id=tid, headers=request.headers
2700
+ )
2701
+ except Exception as exc:
2702
+ if _http_status_code(exc) in {401, 403}:
2703
+ raise
2704
+ await logger.awarning(
2705
+ "Failed to list runs for A2A context",
2706
+ context_id=tid,
2707
+ exc_info=True,
2708
+ )
2709
+ return []
2710
+
2711
+ runs_by_thread = await asyncio.gather(
2712
+ *(list_thread_runs(thread) for thread in threads)
2713
+ )
2714
+
2715
+ # Collect all tasks from matching threads.
2716
+ tasks_by_id: dict[str, dict[str, Any]] = {}
2717
+ seen_task_ids: set[str] = set()
2432
2718
  a2ui_by_task_id: dict[str, A2UIContext] = {}
2433
- for tid in thread_ids:
2434
- try:
2435
- runs = await client.runs.list(
2436
- tid,
2437
- limit=100,
2438
- headers=request.headers,
2719
+ run_id_by_task_id: dict[str, str] = {}
2720
+ task_ids_by_start_id_by_context: dict[str, dict[str, str]] = {}
2721
+ for thread, runs in zip(threads, runs_by_thread, strict=True):
2722
+ tid = str(thread["thread_id"])
2723
+ context_latest_run_id = _latest_run_id(runs)
2724
+ task_ids_by_start_id_by_context[tid] = _task_ids_by_start_id(
2725
+ runs, tid, assistant_id
2726
+ )
2727
+ for run in _runs_for_assistant(runs, assistant_id):
2728
+ task_id = _task_id_for_run(tid, run)
2729
+ if task_id in seen_task_ids:
2730
+ continue
2731
+ seen_task_ids.add(task_id)
2732
+ a2a_state = _task_state_for_run(
2733
+ run,
2734
+ thread if str(run.get("run_id")) == context_latest_run_id else None,
2439
2735
  )
2440
- except Exception:
2441
- continue
2442
-
2443
- for run in runs:
2444
- task_id = _make_task_id(tid, run["run_id"])
2445
- a2a_state = _lg_status_to_a2a_state(run.get("status", "unknown"))
2446
2736
 
2447
2737
  if status_filter and a2a_state != status_filter:
2448
2738
  continue
2449
2739
 
2450
- timestamp = run.get("updated_at") or run.get("created_at") or ""
2451
- if hasattr(timestamp, "isoformat"):
2452
- timestamp = timestamp.isoformat()
2740
+ timestamp = _run_timestamp(run)
2453
2741
 
2454
2742
  if status_timestamp_after and timestamp:
2455
2743
  try:
@@ -2471,10 +2759,12 @@ async def handle_list_tasks(
2471
2759
  if not include_artifacts:
2472
2760
  task["artifacts"] = []
2473
2761
 
2474
- all_tasks.append(task)
2762
+ tasks_by_id[task_id] = task
2475
2763
  a2ui_by_task_id[task_id] = _a2ui_context_for_run(run)
2764
+ run_id_by_task_id[task_id] = str(run["run_id"])
2476
2765
 
2477
2766
  # Sort by timestamp descending (newest first)
2767
+ all_tasks = list(tasks_by_id.values())
2478
2768
  all_tasks.sort(
2479
2769
  key=lambda t: t["status"].get("timestamp", ""),
2480
2770
  reverse=True,
@@ -2487,9 +2777,10 @@ async def handle_list_tasks(
2487
2777
  if history_length != 0:
2488
2778
 
2489
2779
  async def populate_history(task: dict[str, Any]) -> None:
2490
- _, run_id = _parse_task_id(task["id"])
2491
2780
  messages = await _get_historical_messages_for_task(
2492
- task["contextId"], run_id, request.headers
2781
+ task["contextId"],
2782
+ run_id_by_task_id[task["id"]],
2783
+ request.headers,
2493
2784
  )
2494
2785
  task["history"] = _convert_messages_to_a2a_format(
2495
2786
  messages,
@@ -2497,6 +2788,9 @@ async def handle_list_tasks(
2497
2788
  task["contextId"],
2498
2789
  history_length=history_length,
2499
2790
  a2ui=a2ui_by_task_id[task["id"]],
2791
+ task_ids_by_start_id=task_ids_by_start_id_by_context.get(
2792
+ task["contextId"]
2793
+ ),
2500
2794
  )
2501
2795
 
2502
2796
  await asyncio.gather(*(populate_history(task) for task in page_tasks))
@@ -2908,6 +3202,80 @@ async def handle_message_stream(
2908
3202
  )
2909
3203
  return
2910
3204
 
3205
+ context_id_from_message = message.get("contextId")
3206
+ existing_task_id = message.get("taskId")
3207
+ if context_id_from_message is not None and not isinstance(
3208
+ context_id_from_message, str
3209
+ ):
3210
+ yield (
3211
+ b"message",
3212
+ {
3213
+ "jsonrpc": "2.0",
3214
+ "id": rpc_id,
3215
+ "error": {
3216
+ "code": ERROR_CODE_INVALID_PARAMS,
3217
+ "message": "contextId must be a string",
3218
+ },
3219
+ },
3220
+ )
3221
+ return
3222
+ if existing_task_id is not None and not isinstance(existing_task_id, str):
3223
+ yield (
3224
+ b"message",
3225
+ {
3226
+ "jsonrpc": "2.0",
3227
+ "id": rpc_id,
3228
+ "error": {
3229
+ "code": ERROR_CODE_INVALID_PARAMS,
3230
+ "message": "taskId must be a string",
3231
+ },
3232
+ },
3233
+ )
3234
+ return
3235
+
3236
+ is_new_task = existing_task_id is None
3237
+ if existing_task_id is None:
3238
+ context_id = context_id_from_message or str(uuid.uuid4())
3239
+ task_id = None
3240
+ task_start_id = str(uuid7())
3241
+ else:
3242
+ task = await _resolve_task(
3243
+ client,
3244
+ task_id=existing_task_id,
3245
+ context_id=context_id_from_message,
3246
+ assistant_id=assistant_id,
3247
+ headers=request.headers,
3248
+ )
3249
+ if task is None:
3250
+ yield (
3251
+ b"message",
3252
+ {
3253
+ "jsonrpc": "2.0",
3254
+ "id": rpc_id,
3255
+ "error": {
3256
+ "code": ERROR_CODE_TASK_NOT_FOUND,
3257
+ "message": f"Task not found: {existing_task_id}",
3258
+ },
3259
+ },
3260
+ )
3261
+ return
3262
+ if task.state in _TERMINAL_TASK_STATES:
3263
+ yield (
3264
+ b"message",
3265
+ {
3266
+ "jsonrpc": "2.0",
3267
+ "id": rpc_id,
3268
+ "error": {
3269
+ "code": ERROR_CODE_UNSUPPORTED_OPERATION,
3270
+ "message": "A terminal task cannot accept another message",
3271
+ },
3272
+ },
3273
+ )
3274
+ return
3275
+ context_id = task.context_id
3276
+ task_id = task.task_id
3277
+ task_start_id = None
3278
+
2911
3279
  # Process A2A message parts into LangChain messages format
2912
3280
  try:
2913
3281
  a2ui = _a2ui_message_context(message, assistant)
@@ -2917,6 +3285,8 @@ async def handle_message_stream(
2917
3285
  message_role,
2918
3286
  message_id,
2919
3287
  a2ui=a2ui,
3288
+ task_id=task_id,
3289
+ task_start_id=task_start_id,
2920
3290
  )
2921
3291
  except A2UIValidationError as e:
2922
3292
  yield (
@@ -2945,10 +3315,6 @@ async def handle_message_stream(
2945
3315
  )
2946
3316
  return
2947
3317
 
2948
- # Check if this is a continuation (taskId provided in message)
2949
- existing_task_id = message.get("taskId")
2950
- context_id_from_message = message.get("contextId")
2951
-
2952
3318
  # Extract and validate command (LangGraph extension for resuming interrupts)
2953
3319
  command, command_error = _extract_and_validate_command(
2954
3320
  message, context_id_from_message
@@ -2988,7 +3354,9 @@ async def handle_message_stream(
2988
3354
  ) = await _maybe_promote_resume_to_command(
2989
3355
  client=client,
2990
3356
  parts=parts,
2991
- context_id=context_id_from_message,
3357
+ context_id=(
3358
+ context_id if existing_task_id else context_id_from_message
3359
+ ),
2992
3360
  task_id=existing_task_id,
2993
3361
  input_content=input_content,
2994
3362
  headers=request.headers,
@@ -3010,10 +3378,6 @@ async def handle_message_stream(
3010
3378
  task_id=existing_task_id,
3011
3379
  )
3012
3380
 
3013
- if context_id_from_message is None:
3014
- context_id_from_message = str(uuid.uuid4())
3015
- context_id = context_id_from_message
3016
-
3017
3381
  stream = client.runs.stream(
3018
3382
  thread_id=context_id,
3019
3383
  assistant_id=assistant_id,
@@ -3022,7 +3386,12 @@ async def handle_message_stream(
3022
3386
  input=input_content,
3023
3387
  command=command,
3024
3388
  context=run_context,
3025
- metadata=_a2ui_run_metadata(assistant, a2ui),
3389
+ metadata=_a2a_run_metadata(
3390
+ assistant,
3391
+ a2ui,
3392
+ task_id=task_id,
3393
+ task_start_id=task_start_id,
3394
+ ),
3026
3395
  headers=request.headers,
3027
3396
  )
3028
3397
 
@@ -3031,9 +3400,15 @@ async def handle_message_stream(
3031
3400
  run_id = run_event.data.get("run_id")
3032
3401
  if not run_id:
3033
3402
  raise ValueError("Stream did not include run_id")
3403
+ run = await client.runs.get(
3404
+ thread_id=context_id,
3405
+ run_id=str(run_id),
3406
+ headers=request.headers,
3407
+ )
3408
+ context_id = str(run["thread_id"])
3409
+ if is_new_task:
3410
+ task_id = _make_task_id(context_id, str(run_id))
3034
3411
 
3035
- # If continuing an existing task, preserve the original task_id
3036
- task_id = existing_task_id or _make_task_id(context_id, run_id)
3037
3412
  # Emit initial Task object to establish task context
3038
3413
  initial_task = {
3039
3414
  "kind": "task",
langgraph_api/errors.py CHANGED
@@ -59,7 +59,7 @@ async def remote_exception_handler(request, exc: RemoteException):
59
59
  async def overloaded_error_handler(request, exc: ValueError):
60
60
  logger.exception("Overloaded Error", exc_info=exc)
61
61
  return await http_exception_handler(
62
- request, HTTPException(status_code=503, detail=str(exc))
62
+ request, HTTPException(status_code=503, detail="Service Unavailable")
63
63
  )
64
64
 
65
65
 
@@ -97,6 +97,25 @@ def validate_checkpoint_id_from_runnable_config(params: dict[str, Any]) -> str |
97
97
  return None
98
98
 
99
99
 
100
+ def _langsmith_tracer_from_run_start(
101
+ params: dict[str, Any],
102
+ ) -> dict[str, str] | None:
103
+ tracer = params.get("langsmith_tracer")
104
+ if tracer is None:
105
+ return None
106
+ if not _is_record(tracer):
107
+ raise ValueError("langsmith_tracer must be an object.")
108
+ unknown = tracer.keys() - {"project_name", "example_id"}
109
+ if unknown:
110
+ raise ValueError(
111
+ f"langsmith_tracer has unsupported fields: {sorted(unknown)!r}."
112
+ )
113
+ for key, value in tracer.items():
114
+ if not isinstance(value, str):
115
+ raise ValueError(f"langsmith_tracer.{key} must be a string.")
116
+ return dict(tracer)
117
+
118
+
100
119
  # Concurrency strategies accepted on ``run.start`` — the four values the
101
120
  # SDK's ``multitaskStrategy`` option can take, derived from the canonical
102
121
  # ``MultitaskStrategy`` literal so the two never drift.
@@ -481,6 +500,10 @@ class ThreadRunManager:
481
500
  checkpoint_error = validate_checkpoint_id_from_runnable_config(params)
482
501
  if checkpoint_error is not None:
483
502
  return self._error(command.get("id"), "invalid_argument", checkpoint_error)
503
+ try:
504
+ _langsmith_tracer_from_run_start(params)
505
+ except ValueError as exc:
506
+ return self._error(command.get("id"), "invalid_argument", str(exc))
484
507
 
485
508
  try:
486
509
  run = await self._create_or_resume_run(assistant_id, params)
@@ -818,7 +841,7 @@ class ThreadRunManager:
818
841
  # every run.start), falling back to ``enqueue`` — the legacy
819
842
  # stream-endpoint default — when omitted.
820
843
  "multitask_strategy": _multitask_strategy_from_run_start(params),
821
- "langsmith_tracer": None,
844
+ "langsmith_tracer": _langsmith_tracer_from_run_start(params),
822
845
  "durability": None,
823
846
  }
824
847
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: langgraph-api
3
- Version: 0.15.0.dev3
3
+ Version: 0.15.0.dev4
4
4
  Author-email: Will Fu-Hinthorn <will@langchain.dev>, Josh Rogers <josh@langchain.dev>, Parker Rule <parker@langchain.dev>
5
5
  License: Elastic-2.0
6
6
  License-File: LICENSE
@@ -1,4 +1,4 @@
1
- langgraph_api/__init__.py,sha256=nLH816pG3rd8Cu_3NqfysMzBE2nKGNJZWU9-2nJiNfQ,28
1
+ langgraph_api/__init__.py,sha256=xEaLs9f0npDZ1G1mIIH6E14MywcdoOYvZTm8637A30U,28
2
2
  langgraph_api/_factory_utils.py,sha256=_FqW0olMu2PGLsnRq0ylt3a4jDCe12YIFHExMXI6Nyo,5976
3
3
  langgraph_api/asgi_transport.py,sha256=XApY3lIWBZTMbbsl8dDJzl0cLGirmAGE0SifqZUnXvs,11896
4
4
  langgraph_api/asyncio.py,sha256=smqyAoO9nIyPPQRw1IVSCh85fwoZP4PEo-5dDhxGruc,10676
@@ -7,7 +7,7 @@ langgraph_api/cache.py,sha256=LuB3Te0UdXC8a-uEJoRHoe5XsgXulXI59f_Q6DCJNKc,13005
7
7
  langgraph_api/cli.py,sha256=ATtS9s9Cx7QNiGPJceKnMCko29A25ZA-xz39fdxmgfg,22389
8
8
  langgraph_api/command.py,sha256=d-k8h6H4ix1n7fSZ-Zb01NbSkEyqrD6cMKfDFXEIYEw,821
9
9
  langgraph_api/cron_scheduler.py,sha256=dALZwrDn8H3TJFB8S9GzmXiOuGz8_0aUbrXQU38UwZw,6116
10
- langgraph_api/errors.py,sha256=zlMW99wAzNkz2xfik-HMkl_wMqmRFvs1j8V-_DZbAUc,2553
10
+ langgraph_api/errors.py,sha256=i1sO67io5Ffe0T2wlX9s8R1byhf0ZsU5eFX3Pjt2cHU,2566
11
11
  langgraph_api/feature_flags.py,sha256=nPZGxT_alHytycYvOxNi6bp09xUnMpWiHa8LtGmFsO4,2481
12
12
  langgraph_api/graph.py,sha256=ryFwplv0ztnc5g4Sg794GW8xZm2u936uI61fROgU8nQ,38514
13
13
  langgraph_api/http.py,sha256=7hPxKbj-xoAKcm7iucBpT5nM_hXOgGVCPbBsCD693Cw,6977
@@ -41,7 +41,7 @@ langgraph_api/_checkpointer/__init__.py,sha256=ofJTJLGy7Hsuzhj-2dpfDvrDloM0BzlhT
41
41
  langgraph_api/_checkpointer/_adapter.py,sha256=idGYL-FwyVtUXcmhvfYboARfQAkoj0aFD40L5h9N7F8,22289
42
42
  langgraph_api/_checkpointer/protocol.py,sha256=udgYKMNtKWG_eLDwYkHXV3b2bZLZg8Rsfm3fjkhU-rU,3635
43
43
  langgraph_api/api/__init__.py,sha256=88elgNmz6Y1tIi81dzhHa94KVhLvvIqyEJjcov4aH5U,9463
44
- langgraph_api/api/a2a.py,sha256=rvKX9yNM1PC2Zj1zlD6b7agnPX5fq9ylcK4kSPh9gCw,123643
44
+ langgraph_api/api/a2a.py,sha256=l7I6yn_QrrITnUH5-b5MxK663ZjbciRqbk8nMHzFmKU,137203
45
45
  langgraph_api/api/a2ui.py,sha256=rgYWyTYbsBCH9KqRdRzj1T7NhhFVZ_JRlcfR74IyYKo,2952
46
46
  langgraph_api/api/assistants.py,sha256=rwaS3_Sy2nqf_2aIuG5XDvhw0GNWmT1-WIsESMuz1x0,21719
47
47
  langgraph_api/api/event_streaming.py,sha256=jvoBp5qHI92JJHsx4f9LEpb0i0VE31G_qFhS5JQWtlA,17809
@@ -86,7 +86,7 @@ langgraph_api/event_streaming/capabilities.py,sha256=qjVbhCjl1VEQPGeiDxeJAhYGI_7
86
86
  langgraph_api/event_streaming/constants.py,sha256=eGsm-NvOlqV3gNxDO5vlr00FdngmgEQf59zuHSi_74E,1378
87
87
  langgraph_api/event_streaming/event_normalizers.py,sha256=5bVSqGPW-Uh7WX91qgTfwpK433pCSv_wchpbzW8TLi4,2794
88
88
  langgraph_api/event_streaming/namespace.py,sha256=aJDFt45Or2_bQdRpKJgdFhBgTDj6PYl7Coz5GzfbM84,1509
89
- langgraph_api/event_streaming/service.py,sha256=vn0l1Urly4kkFIhHmCWWM2zCu-of9m9wOIq1kPdbJJA,50504
89
+ langgraph_api/event_streaming/service.py,sha256=xraIeJjLBcKkdYw6H10yuR8M8LqB5kC0Xjgq8xmNogw,51363
90
90
  langgraph_api/event_streaming/session.py,sha256=8kx5nxJvJlfIu5fFzQCxEl-8UBJEtxftBLhCYgXMNwo,75947
91
91
  langgraph_api/event_streaming/state_normalizers.py,sha256=vgT4O4tJPr9VDBMn1EP994ieDGDYP43sROnOyLjkEAE,13659
92
92
  langgraph_api/event_streaming/types.py,sha256=RyZqfqgH-jmmmmAFQj5f6nH9M1rGK93zVG7nlmvqZgc,3647
@@ -244,8 +244,8 @@ langgraph_grpc_common/proto/store_pb2.py,sha256=JKcuJK0YUC1BSwLUZJWHCQ7AWkzVS9Kf
244
244
  langgraph_grpc_common/proto/store_pb2.pyi,sha256=ItSuXbA2X9bQlHtuEjDGi0fT_gJj2b0sr7XuQkybe3s,28983
245
245
  langgraph_grpc_common/proto/store_pb2_grpc.py,sha256=Kyt8eV2NXkjh3be_BYytGdD19kNf74961cruonMvCrw,3300
246
246
  langgraph_grpc_common/proto/store_pb2_grpc.pyi,sha256=oS1h2cDoq2OjyM8xazcmc8LwRDC5oCAD0zto2QUmPQw,2027
247
- langgraph_api-0.15.0.dev3.dist-info/METADATA,sha256=psnch_n-wl5jpxPICdPx33XTY1R_YmZE8xOmGdY_BIE,4630
248
- langgraph_api-0.15.0.dev3.dist-info/WHEEL,sha256=zOwg4jB6zX2kU910N-cMawjivD6tO8NEWvE12je1bVk,87
249
- langgraph_api-0.15.0.dev3.dist-info/entry_points.txt,sha256=hGedv8n7cgi41PypMfinwS_HfCwA7xJIfS0jAp8htV8,78
250
- langgraph_api-0.15.0.dev3.dist-info/licenses/LICENSE,sha256=ZPwVR73Biwm3sK6vR54djCrhaRiM4cAD2zvOQZV8Xis,3859
251
- langgraph_api-0.15.0.dev3.dist-info/RECORD,,
247
+ langgraph_api-0.15.0.dev4.dist-info/METADATA,sha256=ou5TufTIWr9twyRuUdpwBmLf2MWsvP3e4vWSMw2qSH4,4630
248
+ langgraph_api-0.15.0.dev4.dist-info/WHEEL,sha256=zOwg4jB6zX2kU910N-cMawjivD6tO8NEWvE12je1bVk,87
249
+ langgraph_api-0.15.0.dev4.dist-info/entry_points.txt,sha256=hGedv8n7cgi41PypMfinwS_HfCwA7xJIfS0jAp8htV8,78
250
+ langgraph_api-0.15.0.dev4.dist-info/licenses/LICENSE,sha256=ZPwVR73Biwm3sK6vR54djCrhaRiM4cAD2zvOQZV8Xis,3859
251
+ langgraph_api-0.15.0.dev4.dist-info/RECORD,,